From 638359aa9d6868dc0c7f5919e0596490311e8ee8 Mon Sep 17 00:00:00 2001
From: alterego655 <824662526@qq.com>
Date: Thu, 9 Jul 2026 20:13:24 +0800
Subject: [PATCH v1] 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 | 33 +++++++++++++++++++++++++++
 src/backend/storage/lmgr/proc.c       |  6 -----
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/src/backend/access/transam/xlogwait.c b/src/backend/access/transam/xlogwait.c
index 582dde3b061..aa6e549ea63 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,32 @@ 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 the process-exit cleanup callback before this process can enter a
+ * wait-LSN heap.  One backend can execute WAIT FOR LSN more than once, so
+ * remember whether the callback is already registered.
+ */
+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 +443,8 @@ WaitForLSN(WaitLSNType lsnType, XLogRecPtr targetLSN, int64 timeout)
 	/* Should have a valid proc number */
 	Assert(MyProcNumber >= 0 && MyProcNumber < MaxBackends + NUM_AUXILIARY_PROCS);
 
+	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 59640bb4f97..2588b14826e 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"
@@ -963,11 +962,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