From b48718e36bbb93b887cb3a7e1ce7d9606466c6eb Mon Sep 17 00:00:00 2001
From: alterego655 <824662526@qq.com>
Date: Fri, 24 Jul 2026 15:44:16 +0800
Subject: [PATCH v6 1/5] Add wait-for-lsn process-exit cleanup callback

Register an on_shmem_exit callback lazily before a process enters a
wait-for-lsn heap, so stale wait state is removed if the process exits while
waiting.  This keeps cleanup local to xlogwait.c and covers non-backend
callers as well.
---
 src/backend/access/transam/xlogwait.c | 38 +++++++++++++++++++++++++++
 src/backend/storage/lmgr/proc.c       |  6 -----
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/src/backend/access/transam/xlogwait.c b/src/backend/access/transam/xlogwait.c
index 582dde3b061..85de3f5cdae 100644
--- a/src/backend/access/transam/xlogwait.c
+++ b/src/backend/access/transam/xlogwait.c
@@ -54,6 +54,7 @@
 #include "miscadmin.h"
 #include "pgstat.h"
 #include "replication/walreceiver.h"
+#include "storage/ipc.h"
 #include "storage/latch.h"
 #include "storage/proc.h"
 #include "storage/shmem.h"
@@ -69,8 +70,12 @@ static int	waitlsn_cmp(const pairingheap_node *a, const pairingheap_node *b,
 
 struct WaitLSNState *waitLSNState = NULL;
 
+static bool waitLSNShmemExitRegistered = false;
+
 static void WaitLSNShmemRequest(void *arg);
 static void WaitLSNShmemInit(void *arg);
+static void WaitLSNShmemExit(int code, Datum arg);
+static void RegisterWaitLSNShmemExit(void);
 
 const ShmemCallbacks WaitLSNShmemCallbacks = {
 	.request_fn = WaitLSNShmemRequest,
@@ -378,6 +383,31 @@ WaitLSNCleanup(void)
 	}
 }
 
+/*
+ * Exit callback to clean up any LSN wait state left behind if this process
+ * exits while waiting.  Transaction abort paths call WaitLSNCleanup()
+ * directly.
+ */
+static void
+WaitLSNShmemExit(int code, Datum arg)
+{
+	WaitLSNCleanup();
+}
+
+/*
+ * Register shared-memory exit cleanup once per process.  A backend may
+ * execute WAIT FOR LSN more than once.
+ */
+static void
+RegisterWaitLSNShmemExit(void)
+{
+	if (!waitLSNShmemExitRegistered)
+	{
+		on_shmem_exit(WaitLSNShmemExit, 0);
+		waitLSNShmemExitRegistered = true;
+	}
+}
+
 /*
  * Check if the given LSN type requires recovery to be in progress.
  * Standby wait types (replay, write, flush) require recovery;
@@ -412,6 +442,14 @@ WaitForLSN(WaitLSNType lsnType, XLogRecPtr targetLSN, int64 timeout)
 	/* Should have a valid proc number */
 	Assert(MyProcNumber >= 0 && MyProcNumber < MaxBackends + NUM_AUXILIARY_PROCS);
 
+	/*
+	 * Ensure cleanup is registered before publishing our waiter entry.
+	 * on_shmem_exit callbacks run in reverse registration order, so this
+	 * callback runs before the earlier-registered ProcKill() and removes the
+	 * entry before our PGPROC slot can be reused.
+	 */
+	RegisterWaitLSNShmemExit();
+
 	if (timeout > 0)
 	{
 		endtime = TimestampTzPlusMilliseconds(GetCurrentTimestamp(), timeout);
diff --git a/src/backend/storage/lmgr/proc.c b/src/backend/storage/lmgr/proc.c
index 9d6e69175a5..624cecc6bc5 100644
--- a/src/backend/storage/lmgr/proc.c
+++ b/src/backend/storage/lmgr/proc.c
@@ -37,7 +37,6 @@
 #include "access/transam.h"
 #include "access/twophase.h"
 #include "access/xlogutils.h"
-#include "access/xlogwait.h"
 #include "miscadmin.h"
 #include "pgstat.h"
 #include "postmaster/autovacuum.h"
@@ -965,11 +964,6 @@ ProcKill(int code, Datum arg)
 	 */
 	LWLockReleaseAll();
 
-	/*
-	 * Cleanup waiting for LSN if any.
-	 */
-	WaitLSNCleanup();
-
 	/* Cancel any pending condition variable sleep, too */
 	ConditionVariableCancelSleep();
 
-- 
2.51.0

