From 96a5df8ed096a117276c89da1968139e1d20c997 Mon Sep 17 00:00:00 2001
From: Alexander Korotkov <akorotkov@postgresql.org>
Date: Sun, 3 Nov 2024 22:47:05 +0200
Subject: [PATCH v1] Teach pg_wal_replay_wait() to handle REPEATABLE READ
 default level

Reported-by:
Bug:
Discussion:
Author:
Reviewed-by:
Tested-by:
Backpatch-through:
---
 src/backend/access/transam/xlogfuncs.c     | 24 ++++++++++++++++++++--
 src/test/recovery/t/043_wal_replay_wait.pl |  5 ++++-
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/src/backend/access/transam/xlogfuncs.c b/src/backend/access/transam/xlogfuncs.c
index bca1d395683..8e90e6ae444 100644
--- a/src/backend/access/transam/xlogfuncs.c
+++ b/src/backend/access/transam/xlogfuncs.c
@@ -34,6 +34,7 @@
 #include "utils/builtins.h"
 #include "utils/memutils.h"
 #include "utils/pg_lsn.h"
+#include "utils/portal.h"
 #include "utils/snapmgr.h"
 #include "utils/timestamp.h"
 
@@ -763,6 +764,8 @@ pg_wal_replay_wait(PG_FUNCTION_ARGS)
 	XLogRecPtr	target_lsn = PG_GETARG_LSN(0);
 	int64		timeout = PG_GETARG_INT64(1);
 	bool		no_error = PG_GETARG_BOOL(2);
+	CallContext *context = (CallContext *) fcinfo->context;
+	bool		start_tranaction = false;
 
 	if (timeout < 0)
 		ereport(ERROR,
@@ -776,12 +779,26 @@ pg_wal_replay_wait(PG_FUNCTION_ARGS)
 	 * implying a kind of self-deadlock.  This is the reason why
 	 * pg_wal_replay_wait() is a procedure, not a function.
 	 *
-	 * At first, we should check there is no active snapshot.  According to
+	 * If we're in REPEATABLE READ isolation level or higher, the catalog
+	 * snapshot used to find this procedure will be saved as the transaction
+	 * snapshot.  Thus, we need to finish current transaction in order to
+	 * release this snapshot.  We can do this in non-atomic context.
+	 * Additionally check we're in the first command of the transaction.
+	 *
+	 * Also, we should check there is no active snapshot.  According to
 	 * PlannedStmtRequiresSnapshot(), even in an atomic context, CallStmt is
 	 * processed with a snapshot.  Thankfully, we can pop this snapshot,
 	 * because PortalRunUtility() can tolerate this.
 	 */
-	if (ActiveSnapshotSet())
+	if (IsolationUsesXactSnapshot() &&
+		GetCurrentCommandId(false) == 0 &&
+		!context->atomic)
+	{
+		ForgetPortalSnapshots();
+		CommitTransactionCommand();
+		start_tranaction = true;
+	}
+	else if (ActiveSnapshotSet())
 		PopActiveSnapshot();
 
 	/*
@@ -805,6 +822,9 @@ pg_wal_replay_wait(PG_FUNCTION_ARGS)
 
 	lastWaitLSNResult = WaitForLSNReplay(target_lsn, timeout);
 
+	if (start_tranaction)
+		StartTransactionCommand();
+
 	if (no_error)
 		PG_RETURN_VOID();
 
diff --git a/src/test/recovery/t/043_wal_replay_wait.pl b/src/test/recovery/t/043_wal_replay_wait.pl
index 5857b943711..388741ee0aa 100644
--- a/src/test/recovery/t/043_wal_replay_wait.pl
+++ b/src/test/recovery/t/043_wal_replay_wait.pl
@@ -47,13 +47,16 @@ my $output = $node_standby->safe_psql(
 ok($output >= 0,
 	"standby reached the same LSN as primary after pg_wal_replay_wait()");
 
-# 2. Check that new data is visible after calling pg_wal_replay_wait()
+# 2. Check that new data is visible after calling pg_wal_replay_wait().
+# Do this check with REPEATABLE READ as the default transaction isolation
+# mode.
 $node_primary->safe_psql('postgres',
 	"INSERT INTO wait_test VALUES (generate_series(21, 30))");
 my $lsn2 =
   $node_primary->safe_psql('postgres', "SELECT pg_current_wal_insert_lsn()");
 $output = $node_standby->safe_psql(
 	'postgres', qq[
+	SET default_transaction_isolation = 'repeatable read';
 	CALL pg_wal_replay_wait('${lsn2}');
 	SELECT count(*) FROM wait_test;
 ]);
-- 
2.39.5 (Apple Git-154)

