From 200a46caf27e6f02bb58ac107d8d10dfc47436d4 Mon Sep 17 00:00:00 2001
From: alterego655 <824662526@qq.com>
Date: Mon, 27 Jul 2026 18:02:45 +0800
Subject: [PATCH v1 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         |  27 +++++
 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      |   8 ++
 .../modules/test_wait_lsn/test_wait_lsn.c     |  36 ++++++
 .../test_wait_lsn/test_wait_lsn.control       |   4 +
 src/test/recovery/Makefile                    |   3 +-
 src/test/recovery/t/049_wait_for_lsn.pl       | 104 ++++++++++++++++++
 10 files changed, 226 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..0f30aae4660 100644
--- a/src/backend/access/transam/xlogwait.c
+++ b/src/backend/access/transam/xlogwait.c
@@ -60,6 +60,7 @@
 #include "storage/shmem.h"
 #include "storage/subsystems.h"
 #include "utils/fmgrprotos.h"
+#include "utils/injection_point.h"
 #include "utils/pg_lsn.h"
 #include "utils/snapmgr.h"
 #include "utils/wait_event.h"
@@ -239,6 +240,12 @@ addLSNWaiter(XLogRecPtr lsn, WaitLSNType lsnType)
 	updateMinWaitedLSN(lsnType);
 
 	LWLockRelease(WaitLSNLock);
+
+	/*
+	 * Note the placement outside of the lock: this injection point may block,
+	 * which would stall every other registration and wakeup.
+	 */
+	INJECTION_POINT("wait-for-lsn-after-register", NULL);
 }
 
 /*
@@ -437,6 +444,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 +454,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 +499,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..d830fd97b95
--- /dev/null
+++ b/src/test/modules/test_wait_lsn/test_wait_lsn--1.0.sql
@@ -0,0 +1,8 @@
+/* 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;
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..82752005597
--- /dev/null
+++ b/src/test/modules/test_wait_lsn/test_wait_lsn.c
@@ -0,0 +1,36 @@
+/*--------------------------------------------------------------------------
+ *
+ * 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 "utils/pg_lsn.h"
+
+PG_MODULE_MAGIC;
+
+PG_FUNCTION_INFO_V1(test_wait_lsn_stale_wakeup);
+
+/*
+ * Wake standby_write waiters through the supplied LSN without changing the
+ * actual write or replay positions.  This simulates a notification made stale
+ * by a subsequent decrease in the walreceiver's published write position.
+ */
+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();
+}
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..2405c2ef5f8 100644
--- a/src/test/recovery/t/049_wait_for_lsn.pl
+++ b/src/test/recovery/t/049_wait_for_lsn.pl
@@ -1055,6 +1055,110 @@ 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 'Injection points not supported by this build', 2
+	  unless $ENV{enable_injection_points} eq 'yes';
+	skip 'Required test extensions are not installed', 2
+	  unless $rcv_primary->check_extension('injection_points')
+	  && $rcv_primary->check_extension('test_wait_lsn');
+
+	my $register_point = 'wait-for-lsn-after-register';
+	my $stale_appname = 'standby_write_stale_wakeup';
+
+	$rcv_primary->safe_psql(
+		'postgres', q[
+		CREATE EXTENSION injection_points;
+		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()');
+
+	# Set application_name in the connection string so that the waiter can be
+	# identified without issuing a preliminary SQL command.
+	my $stale_session = $rcv_standby->background_psql(
+		'postgres',
+		connstr => $rcv_standby->connstr('postgres')
+		  . " application_name=$stale_appname");
+	$stale_session->set_query_timer_restart();
+	$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 count(*) = 1
+		FROM pg_stat_activity
+		WHERE application_name = '$stale_appname'
+		  AND wait_event = 'WaitForWalWrite'
+	]
+	) or die "standby_write waiter did not sleep";
+
+	# Waiting in WaitForWalWrite proves that the initial registration is
+	# complete.  Attach only now so that this waiter can reach the injection
+	# point only by re-registering.
+	$rcv_standby->safe_psql('postgres',
+		"SELECT injection_points_attach('$register_point', 'wait')");
+
+	# 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')");
+
+	# Injection points are global, so use application_name to verify that this
+	# waiter, rather than another backend, reached the registration point.
+	# An unrelated latch wake cannot satisfy this condition because only a
+	# call to addLSNWaiter() can reach the point.
+	ok( $rcv_standby->poll_query_until(
+			'postgres', qq[
+		SELECT count(*) = 1
+		FROM pg_stat_activity
+		WHERE application_name = '$stale_appname'
+		  AND wait_event = '$register_point'
+	]),
+		"standby_write waiter re-registers after a stale wakeup")
+	  or die "standby_write waiter did not re-register";
+
+	$rcv_standby->safe_psql(
+		'postgres', qq[
+		SELECT injection_points_detach('$register_point');
+		SELECT injection_points_wakeup('$register_point');
+	]);
+
+	# Resume streaming so that the target is reached for real.  Deliberately
+	# omit a WAIT FOR timeout: it would make WaitLatch() return and recheck the
+	# position, potentially reporting success even if the WAL-progress
+	# notification was lost.
+	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

