From bad61e8a046c15a8afb365767034d76f9f6ceb68 Mon Sep 17 00:00:00 2001
From: alterego655 <824662526@qq.com>
Date: Mon, 27 Jul 2026 18:02:45 +0800
Subject: [PATCH v2 5/5] Re-register LSN waiters after stale wakeups

WaitLSNWakeup() removes a waiter from the heap before setting its latch.
If the position that caused the wakeup moves backwards before the waiter
rechecks it, as can happen when WAL streaming restarts, the waiter may
sleep again while no longer registered.  Subsequent WAL progress then
cannot wake it.

When an unmet waiter finds that it is no longer in the heap, add it back
and restart the loop.  Rereading the position after registration also
prevents missing an advance between the previous read and the re-add.

Add deterministic TAP coverage that simulates a stale standby_write
wakeup without advancing the actual write or replay positions.
---
 src/backend/access/transam/xlogwait.c         |  20 ++++
 src/test/modules/Makefile                     |   1 +
 src/test/modules/meson.build                  |   1 +
 src/test/modules/test_wait_lsn/Makefile       |  21 ++++
 src/test/modules/test_wait_lsn/meson.build    |  22 ++++
 .../test_wait_lsn/test_wait_lsn--1.0.sql      |  13 +++
 .../modules/test_wait_lsn/test_wait_lsn.c     | 101 ++++++++++++++++++
 .../test_wait_lsn/test_wait_lsn.control       |   4 +
 src/test/recovery/Makefile                    |   3 +-
 src/test/recovery/t/049_wait_for_lsn.pl       |  85 +++++++++++++++
 10 files changed, 270 insertions(+), 1 deletion(-)
 create mode 100644 src/test/modules/test_wait_lsn/Makefile
 create mode 100644 src/test/modules/test_wait_lsn/meson.build
 create mode 100644 src/test/modules/test_wait_lsn/test_wait_lsn--1.0.sql
 create mode 100644 src/test/modules/test_wait_lsn/test_wait_lsn.c
 create mode 100644 src/test/modules/test_wait_lsn/test_wait_lsn.control

diff --git a/src/backend/access/transam/xlogwait.c b/src/backend/access/transam/xlogwait.c
index f60491b67d1..a52ad25a64e 100644
--- a/src/backend/access/transam/xlogwait.c
+++ b/src/backend/access/transam/xlogwait.c
@@ -437,6 +437,7 @@ WaitLSNResult
 WaitForLSN(WaitLSNType lsnType, XLogRecPtr targetLSN, int64 timeout)
 {
 	XLogRecPtr	currentLSN;
+	WaitLSNProcInfo *procInfo;
 	TimestampTz endtime = 0;
 	int			wake_events = WL_LATCH_SET | WL_POSTMASTER_DEATH;
 
@@ -446,6 +447,8 @@ WaitForLSN(WaitLSNType lsnType, XLogRecPtr targetLSN, int64 timeout)
 	/* Should have a valid proc number */
 	Assert(MyProcNumber >= 0 && MyProcNumber < MaxBackends + NUM_AUXILIARY_PROCS);
 
+	procInfo = &waitLSNState->procInfos[MyProcNumber];
+
 	RegisterWaitLSNShmemExit();
 
 	if (timeout > 0)
@@ -489,6 +492,23 @@ WaitForLSN(WaitLSNType lsnType, XLogRecPtr targetLSN, int64 timeout)
 				break;
 		}
 
+		/*
+		 * The target is not reached.  Normally we remain in the waiters heap
+		 * and can sleep again.  A wakeup can become stale, however, if the
+		 * position moves backwards after the waker removes us (for example,
+		 * after a walreceiver restart).  Re-register in that case and reread
+		 * the position, since an advance before the re-add could not wake us.
+		 *
+		 * It is safe to read inHeap without the lock because only this process
+		 * sets it true.  If a waker clears it concurrently, it also sets our
+		 * latch, so we will recheck and re-register if necessary.
+		 */
+		if (!procInfo->inHeap)
+		{
+			addLSNWaiter(targetLSN, lsnType);
+			continue;
+		}
+
 		if (timeout > 0)
 		{
 			delay_ms = TimestampDifferenceMilliseconds(GetCurrentTimestamp(), endtime);
diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 098bb8142ae..bb88b3058ed 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -53,6 +53,7 @@ SUBDIRS = \
 		  test_shm_mq \
 		  test_slru \
 		  test_tidstore \
+		  test_wait_lsn \
 		  unsafe_tests \
 		  worker_spi \
 		  xid_wraparound
diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build
index 4bca42bb370..ce09e00531d 100644
--- a/src/test/modules/meson.build
+++ b/src/test/modules/meson.build
@@ -54,6 +54,7 @@ subdir('test_shmem')
 subdir('test_shm_mq')
 subdir('test_slru')
 subdir('test_tidstore')
+subdir('test_wait_lsn')
 subdir('typcache')
 subdir('unsafe_tests')
 subdir('worker_spi')
diff --git a/src/test/modules/test_wait_lsn/Makefile b/src/test/modules/test_wait_lsn/Makefile
new file mode 100644
index 00000000000..e9ce2fd3ad0
--- /dev/null
+++ b/src/test/modules/test_wait_lsn/Makefile
@@ -0,0 +1,21 @@
+# src/test/modules/test_wait_lsn/Makefile
+
+MODULE_big = test_wait_lsn
+OBJS = \
+	$(WIN32RES) \
+	test_wait_lsn.o
+PGFILEDESC = "test_wait_lsn - test code for WAIT FOR LSN"
+
+EXTENSION = test_wait_lsn
+DATA = test_wait_lsn--1.0.sql
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_wait_lsn
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_wait_lsn/meson.build b/src/test/modules/test_wait_lsn/meson.build
new file mode 100644
index 00000000000..66e074298d7
--- /dev/null
+++ b/src/test/modules/test_wait_lsn/meson.build
@@ -0,0 +1,22 @@
+# Copyright (c) 2026, PostgreSQL Global Development Group
+
+test_wait_lsn_sources = files(
+  'test_wait_lsn.c',
+)
+
+if host_system == 'windows'
+  test_wait_lsn_sources += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_wait_lsn',
+    '--FILEDESC', 'test_wait_lsn - test code for WAIT FOR LSN',])
+endif
+
+test_wait_lsn = shared_module('test_wait_lsn',
+  test_wait_lsn_sources,
+  kwargs: pg_test_mod_args,
+)
+test_install_libs += test_wait_lsn
+
+test_install_data += files(
+  'test_wait_lsn.control',
+  'test_wait_lsn--1.0.sql',
+)
diff --git a/src/test/modules/test_wait_lsn/test_wait_lsn--1.0.sql b/src/test/modules/test_wait_lsn/test_wait_lsn--1.0.sql
new file mode 100644
index 00000000000..829ef54f73f
--- /dev/null
+++ b/src/test/modules/test_wait_lsn/test_wait_lsn--1.0.sql
@@ -0,0 +1,13 @@
+/* src/test/modules/test_wait_lsn/test_wait_lsn--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_wait_lsn" to load this file. \quit
+
+CREATE FUNCTION test_wait_lsn_stale_wakeup(pg_catalog.pg_lsn)
+RETURNS pg_catalog.void STRICT
+AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION test_wait_lsn_waiter_is_registered(
+    pg_catalog.int4, pg_catalog.text, pg_catalog.pg_lsn)
+RETURNS pg_catalog.bool STRICT
+AS 'MODULE_PATHNAME' LANGUAGE C;
diff --git a/src/test/modules/test_wait_lsn/test_wait_lsn.c b/src/test/modules/test_wait_lsn/test_wait_lsn.c
new file mode 100644
index 00000000000..ac142daa742
--- /dev/null
+++ b/src/test/modules/test_wait_lsn/test_wait_lsn.c
@@ -0,0 +1,101 @@
+/*--------------------------------------------------------------------------
+ *
+ * test_wait_lsn.c
+ *		Test support for WAIT FOR LSN.
+ *
+ * Copyright (c) 2026, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *		src/test/modules/test_wait_lsn/test_wait_lsn.c
+ *
+ * -------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include "access/xlogwait.h"
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "storage/lwlock.h"
+#include "storage/proc.h"
+#include "storage/procarray.h"
+#include "utils/builtins.h"
+#include "utils/pg_lsn.h"
+
+PG_MODULE_MAGIC;
+
+PG_FUNCTION_INFO_V1(test_wait_lsn_stale_wakeup);
+PG_FUNCTION_INFO_V1(test_wait_lsn_waiter_is_registered);
+
+static WaitLSNType
+parse_wait_lsn_type(text *mode_text)
+{
+	char	   *mode = text_to_cstring(mode_text);
+	WaitLSNType lsn_type;
+
+	if (pg_strcasecmp(mode, "standby_replay") == 0)
+		lsn_type = WAIT_LSN_TYPE_STANDBY_REPLAY;
+	else if (pg_strcasecmp(mode, "standby_write") == 0)
+		lsn_type = WAIT_LSN_TYPE_STANDBY_WRITE;
+	else if (pg_strcasecmp(mode, "standby_flush") == 0)
+		lsn_type = WAIT_LSN_TYPE_STANDBY_FLUSH;
+	else if (pg_strcasecmp(mode, "primary_flush") == 0)
+		lsn_type = WAIT_LSN_TYPE_PRIMARY_FLUSH;
+	else
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("unrecognized WAIT FOR LSN mode \"%s\"", mode)));
+
+	pfree(mode);
+	return lsn_type;
+}
+
+/*
+ * Remove standby_write waiters through the supplied LSN from the heap and
+ * wake them without advancing the actual write or replay positions.  This
+ * reproduces the state after a WAL-progress notification becomes stale before
+ * the waiter rechecks its target.
+ */
+Datum
+test_wait_lsn_stale_wakeup(PG_FUNCTION_ARGS)
+{
+	XLogRecPtr	target_lsn = PG_GETARG_LSN(0);
+
+	WaitLSNWakeup(WAIT_LSN_TYPE_STANDBY_WRITE, target_lsn);
+
+	PG_RETURN_VOID();
+}
+
+/*
+ * Check whether the backend with the supplied PID is registered for the
+ * supplied mode and target.  ProcArrayLock stabilizes the PID mapping, while
+ * WaitLSNLock protects the registration state.
+ */
+Datum
+test_wait_lsn_waiter_is_registered(PG_FUNCTION_ARGS)
+{
+	int			pid = PG_GETARG_INT32(0);
+	WaitLSNType lsn_type = parse_wait_lsn_type(PG_GETARG_TEXT_PP(1));
+	XLogRecPtr	target_lsn = PG_GETARG_LSN(2);
+	bool		registered = false;
+	PGPROC	   *proc;
+
+	LWLockAcquire(ProcArrayLock, LW_SHARED);
+	proc = BackendPidGetProcWithLock(pid);
+
+	if (proc != NULL)
+	{
+		ProcNumber	procno = GetNumberFromPGProc(proc);
+		WaitLSNProcInfo *proc_info = &waitLSNState->procInfos[procno];
+
+		LWLockAcquire(WaitLSNLock, LW_SHARED);
+		registered = proc_info->inHeap &&
+			proc_info->procno == procno &&
+			proc_info->lsnType == lsn_type &&
+			proc_info->waitLSN == target_lsn;
+		LWLockRelease(WaitLSNLock);
+	}
+
+	LWLockRelease(ProcArrayLock);
+
+	PG_RETURN_BOOL(registered);
+}
diff --git a/src/test/modules/test_wait_lsn/test_wait_lsn.control b/src/test/modules/test_wait_lsn/test_wait_lsn.control
new file mode 100644
index 00000000000..7b84150e179
--- /dev/null
+++ b/src/test/modules/test_wait_lsn/test_wait_lsn.control
@@ -0,0 +1,4 @@
+comment = 'Test code for WAIT FOR LSN'
+default_version = '1.0'
+module_pathname = '$libdir/test_wait_lsn'
+relocatable = true
diff --git a/src/test/recovery/Makefile b/src/test/recovery/Makefile
index d41aaaf8ae1..9c4102b6b2c 100644
--- a/src/test/recovery/Makefile
+++ b/src/test/recovery/Makefile
@@ -12,7 +12,8 @@
 EXTRA_INSTALL=contrib/pg_prewarm \
 	contrib/pg_stat_statements \
 	contrib/test_decoding \
-	src/test/modules/injection_points
+	src/test/modules/injection_points \
+	src/test/modules/test_wait_lsn
 
 subdir = src/test/recovery
 top_builddir = ../../..
diff --git a/src/test/recovery/t/049_wait_for_lsn.pl b/src/test/recovery/t/049_wait_for_lsn.pl
index bc216064714..826d7826ec2 100644
--- a/src/test/recovery/t/049_wait_for_lsn.pl
+++ b/src/test/recovery/t/049_wait_for_lsn.pl
@@ -1055,6 +1055,91 @@ is($boundary_session->{stdout},
 	'success',
 	"standby_replay: waiter at current + 1 wakes when replay advances");
 
+# 11d. A standby_write waiter removed from the waiters heap by a stale wakeup
+# must re-register if its target has not actually been reached.
+SKIP:
+{
+	skip 'Required test extension is not installed', 2
+	  unless $rcv_primary->check_extension('test_wait_lsn');
+
+	$rcv_primary->safe_psql('postgres', 'CREATE EXTENSION test_wait_lsn');
+	$rcv_primary->wait_for_catchup($rcv_standby);
+
+	# Stop streaming before generating the target.  This keeps both the
+	# walreceiver write position and the replay position below the target.
+	stop_walreceiver($rcv_standby);
+	$rcv_primary->safe_psql('postgres', 'INSERT INTO rcv_test VALUES (300)');
+	my $stale_target = $rcv_primary->safe_psql('postgres',
+		'SELECT pg_current_wal_insert_lsn()');
+
+	# Keep WAIT FOR untimed.  After streaming resumes below, its own timeout
+	# could wake the backend and make the completion check pass even if the
+	# WAL-progress wakeup were lost.
+	my $stale_waiter_name = 'wait_for_lsn_stale_wakeup';
+	my $stale_session = $rcv_standby->background_psql('postgres',
+		connstr => $rcv_standby->connstr('postgres')
+		  . " application_name=$stale_waiter_name");
+
+	$stale_session->set_query_timer_restart();
+
+	my $stale_waiter_pid = $rcv_standby->safe_psql(
+		'postgres',
+		"SELECT pid FROM pg_stat_activity
+		 WHERE application_name = '$stale_waiter_name'");
+	die "could not determine stale waiter PID: '$stale_waiter_pid'"
+	  unless $stale_waiter_pid =~ /^[0-9]+$/;
+
+	$stale_session->query_until(
+		qr/started/, qq[
+		\\echo started
+		WAIT FOR LSN '$stale_target' WITH (MODE 'standby_write');
+		\\echo completed
+	]);
+
+	$rcv_standby->poll_query_until(
+		'postgres', qq[
+		SELECT test_wait_lsn_waiter_is_registered(
+			$stale_waiter_pid, 'standby_write', '$stale_target')
+	]
+	) or die "standby_write waiter did not register";
+
+	# The write position is below the target because the receiver was stopped
+	# before the target was generated.  Verify that replay is below it too.
+	$rcv_standby->safe_psql(
+		'postgres', qq[
+		SELECT pg_wal_lsn_diff(
+			'$stale_target'::pg_lsn, pg_last_wal_replay_lsn()) > 0
+	]) eq 't'
+	  or die "standby replay reached the target before the stale wakeup";
+
+	# Simulate the state left by a wakeup that became stale before the waiter
+	# could recheck its target.  The helper removes the waiter from the heap
+	# and wakes it without changing the actual write or replay positions.  To
+	# the waiter, this is indistinguishable from the position reaching the
+	# target and then decreasing before the recheck.
+	$rcv_standby->safe_psql('postgres',
+		"SELECT test_wait_lsn_stale_wakeup('$stale_target')");
+
+	# The position is still below the target, so the waiter must restore its
+	# heap registration before sleeping again.
+	ok( $rcv_standby->poll_query_until(
+			'postgres', qq[
+		SELECT test_wait_lsn_waiter_is_registered(
+			$stale_waiter_pid, 'standby_write', '$stale_target')
+	]),
+		"standby_write waiter re-registers after a stale wakeup"
+	) or die "standby_write waiter did not re-register";
+
+	# Reach the target for real; WAL progress should wake the re-registered waiter.
+	resume_walreceiver($rcv_standby);
+
+	like(
+		$stale_session->query_until(qr/completed/, ''),
+		qr/^success\r?\ncompleted/m,
+		"standby_write waiter completes once the target is reached");
+	$stale_session->quit;
+}
+
 $rcv_standby->stop;
 $rcv_primary->stop;
 
-- 
2.51.0

