Add wait-for-lsn process-exit cleanup callback

WaitLSNCleanup() was called from ProcKill(), but not from
AuxiliaryProcKill(), even though xlogwait.c sizes its shared memory to
include NUM_AUXILIARY_PROCS and thus accepts calls from auxiliary
processes.  Such a process exiting while waiting would leave its entry in
the heap.

Register an on_shmem_exit callback lazily before a process enters a
wait-for-lsn heap instead, as suggested by Noah: that keeps the cleanup
local to xlogwait.c and makes it harder to miss a caller that needs it,
rather than having to remember every process-kill path.

Reported-by: Noah Misch <noah@leadboat.com>
Author: Xuneng Zhou <xunengzhou@gmail.com>
Reviewed-by: Alexander Korotkov <aekorotkov@gmail.com>
Discussion: https://postgr.es/m/20260706012642.f9.noahmisch%40microsoft.com
Discussion: https://postgr.es/m/CABPTF7UtW_cAa%3DQh4RDfKiUqu3pJJE22ai9tbWJVERbeRyssLw%40mail.gmail.com
Backpatch-through: 19

diff --git a/src/backend/access/transam/xlogwait.c b/src/backend/access/transam/xlogwait.c
index 582dde3b061281947dd9d3ff9eb66de63a6821cf..85de3f5cdaea59e29954aa2c14c3d25d3aa51c68 100644 (file)
--- 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 9d6e69175a58a70240f00c07ede073e8bf50cbf6..624cecc6bc5f6a3f1779032a4bb06f6a9fa9d014 100644 (file)
--- 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();