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.

Process interrupts before re-registering rather than after, so that a
wakeup which goes stale again cannot postpone cancellation, however often
it repeats.  As a side effect, a pending cancel now wins over an expired
timeout, which previously reported a timeout instead.

Add deterministic TAP coverage that simulates a stale standby_write
wakeup without advancing the actual write or replay positions.

Author: Xuneng Zhou <xunengzhou@gmail.com>
Reviewed-by: Alexander Korotkov <aekorotkov@gmail.com>
Discussion: https://postgr.es/m/CABPTF7UtW_cAa%3DQh4RDfKiUqu3pJJE22ai9tbWJVERbeRyssLw%40mail.gmail.com
Backpatch-through: 19

src/backend/access/transam/xlogwait.c

@@ -440,6 +440,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;
 
@@ -449,6 +450,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];
+
 	/*
 	 * Ensure cleanup is registered before publishing our waiter entry.
 	 * on_shmem_exit callbacks run in reverse registration order, so this
@@ -498,6 +501,38 @@ WaitForLSN(WaitLSNType lsnType, XLogRecPtr targetLSN, int64 timeout)
 				break;
 		}
 
+		CHECK_FOR_INTERRUPTS();
+
+		/*
+		 * 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 removed us.  That happens
+		 * with the walreceiver-tracked positions: when streaming starts on a
+		 * new timeline, or after receiveStart was reset,
+		 * RequestXLogStreaming() re-seeds writtenUpto and flushedUpto with
+		 * the requested start position, which can be below what was published
+		 * before.  Re-register in that case and reread the position, since an
+		 * advance between the previous read and the re-add could not have
+		 * woken us.
+		 *
+		 * A wakeup that goes stale again sends us around the loop once more,
+		 * so interrupts are processed before we re-register: however often
+		 * that repeats, the wait stays cancellable.  Repeating requires a
+		 * fresh wakeup, hence the position reaching the target and falling
+		 * back below it, so it follows streaming restarts rather than burning
+		 * CPU.  The deadline is checked on every iteration that goes on to
+		 * sleep, which is the only place it matters.
+		 *
+		 * 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);
@@ -505,8 +540,6 @@ WaitForLSN(WaitLSNType lsnType, XLogRecPtr targetLSN, int64 timeout)
 				break;
 		}
 
-		CHECK_FOR_INTERRUPTS();
-
 		rc = WaitLatch(MyLatch, wake_events, delay_ms,
 					   WaitLSNWaitEvents[lsnType]);
 

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

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')

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

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',
+)

src/test/modules/test_wait_lsn/test_wait_lsn--1.0.sql

@@ -0,0 +1,14 @@
+/* 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_wakeup(
+    pg_catalog.text, 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;

src/test/modules/test_wait_lsn/test_wait_lsn.c

@@ -0,0 +1,99 @@
+/*--------------------------------------------------------------------------
+ *
+ * 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 "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_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;
+}
+
+/*
+ * Wake all waiters of the supplied type through the supplied LSN without
+ * advancing the underlying WAL position.
+ */
+Datum
+test_wait_lsn_wakeup(PG_FUNCTION_ARGS)
+{
+	WaitLSNType lsn_type = parse_wait_lsn_type(PG_GETARG_TEXT_PP(0));
+	XLogRecPtr	upto_lsn = PG_GETARG_LSN(1);
+
+	WaitLSNWakeup(lsn_type, upto_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);
+}

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

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 = ../../..

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_wakeup('standby_write', '$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;