v2-0002-Implement-pg_wal_replay_wait_no_error.patch

From af3d86506e5dd20b84c05c54f66d57fb480f1904 Mon Sep 17 00:00:00 2001
From: Alexander Korotkov <akorotkov@postgresql.org>
Date: Thu, 19 Sep 2024 15:34:18 +0300
Subject: [PATCH v2 2/2] Implement pg_wal_replay_wait_no_error()

And function pg_wal_replay_wait_status() to get the last status.
---
 src/backend/access/transam/xlogfuncs.c     | 93 ++++++++++++++++++----
 src/backend/catalog/system_functions.sql   |  5 ++
 src/include/catalog/pg_proc.dat            | 11 +++
 src/test/recovery/t/043_wal_replay_wait.pl | 12 +++
 4 files changed, 105 insertions(+), 16 deletions(-)

diff --git a/src/backend/access/transam/xlogfuncs.c b/src/backend/access/transam/xlogfuncs.c
index 74b493e437e..6602507f452 100644
--- a/src/backend/access/transam/xlogfuncs.c
+++ b/src/backend/access/transam/xlogfuncs.c
@@ -752,20 +752,11 @@ pg_promote(PG_FUNCTION_ARGS)
 }
 
 /*
- * Waits until recovery replays the target LSN with optional timeout.
+ * Prepare for waiting for LSN replay.
  */
-Datum
-pg_wal_replay_wait(PG_FUNCTION_ARGS)
+static void
+pg_wal_replay_wait_prepare(const char *funcname)
 {
-	XLogRecPtr	target_lsn = PG_GETARG_LSN(0);
-	int64		timeout = PG_GETARG_INT64(1);
-	WaitLSNResult result;
-
-	if (timeout < 0)
-		ereport(ERROR,
-				(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
-				 errmsg("\"timeout\" must not be negative")));
-
 	/*
 	 * We are going to wait for the LSN replay.  We should first care that we
 	 * don't hold a snapshot and correspondingly our MyProc->xmin is invalid.
@@ -791,22 +782,42 @@ pg_wal_replay_wait(PG_FUNCTION_ARGS)
 	if (GetOldestSnapshot())
 		ereport(ERROR,
 				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
-				 errmsg("pg_wal_replay_wait() must be only called without an active or registered snapshot"),
-				 errdetail("Make sure pg_wal_replay_wait() isn't called within a transaction with an isolation level higher than READ COMMITTED, another procedure, or a function.")));
+				 errmsg("%s must be only called without an active or registered snapshot", funcname),
+				 errdetail("Make sure %s isn't called within a transaction with an isolation level higher than READ COMMITTED, another procedure, or a function.", funcname)));
 
 	/*
 	 * As the result we should hold no snapshot, and correspondingly our xmin
 	 * should be unset.
 	 */
 	Assert(MyProc->xmin == InvalidTransactionId);
+}
 
-	result = WaitForLSNReplay(target_lsn, timeout);
+static WaitLSNResult lastWaitLSNResult = WaitLSNResultSuccess;
+
+/*
+ * Waits until recovery replays the target LSN with optional timeout.  Throw
+ * an error on failure.
+ */
+Datum
+pg_wal_replay_wait(PG_FUNCTION_ARGS)
+{
+	XLogRecPtr	target_lsn = PG_GETARG_LSN(0);
+	int64		timeout = PG_GETARG_INT64(1);
+
+	if (timeout < 0)
+		ereport(ERROR,
+				(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
+				 errmsg("\"timeout\" must not be negative")));
+
+	pg_wal_replay_wait_prepare("pg_wal_replay_wait()");
+
+	lastWaitLSNResult = WaitForLSNReplay(target_lsn, timeout);
 
 	/*
 	 * Process the result of WaitForLSNReplay().  Throw appropriate error if
 	 * needed.
 	 */
-	switch (result)
+	switch (lastWaitLSNResult)
 	{
 		case WaitLSNResultSuccess:
 			/* Nothing to do on success */
@@ -839,3 +850,53 @@ pg_wal_replay_wait(PG_FUNCTION_ARGS)
 
 	PG_RETURN_VOID();
 }
+
+/*
+ * Waits until recovery replays the target LSN with optional timeout.  Return
+ * the waiting result as a text.
+ */
+Datum
+pg_wal_replay_wait_no_error(PG_FUNCTION_ARGS)
+{
+	XLogRecPtr	target_lsn = PG_GETARG_LSN(0);
+	int64		timeout = PG_GETARG_INT64(1);
+
+	if (timeout < 0)
+		ereport(ERROR,
+				(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
+				 errmsg("\"timeout\" must not be negative")));
+
+	pg_wal_replay_wait_prepare("pg_wal_replay_wait_status()");
+
+	lastWaitLSNResult = WaitForLSNReplay(target_lsn, timeout);
+
+	return (Datum) 0;
+}
+
+Datum
+pg_wal_replay_wait_status(PG_FUNCTION_ARGS)
+{
+	const char *result_string = "";
+
+	/* Process the result of WaitForLSNReplay(). */
+	switch (lastWaitLSNResult)
+	{
+		case WaitLSNResultSuccess:
+			result_string = "success";
+			break;
+
+		case WaitLSNResultTimeout:
+			result_string = "timeout";
+			break;
+
+		case WaitLSNResultNotInRecovery:
+			result_string = "not in recovery";
+			break;
+
+		case WaitLSNResultPromotedConcurrently:
+			result_string = "promoted concurrently";
+			break;
+	}
+
+	PG_RETURN_TEXT_P(cstring_to_text(result_string));
+}
diff --git a/src/backend/catalog/system_functions.sql b/src/backend/catalog/system_functions.sql
index b0d0de051e7..ed092b748ef 100644
--- a/src/backend/catalog/system_functions.sql
+++ b/src/backend/catalog/system_functions.sql
@@ -417,6 +417,11 @@ CREATE OR REPLACE FUNCTION
 CREATE OR REPLACE PROCEDURE pg_wal_replay_wait(target_lsn pg_lsn, timeout int8 DEFAULT 0)
   LANGUAGE internal AS 'pg_wal_replay_wait';
 
+CREATE OR REPLACE PROCEDURE pg_wal_replay_wait_status(OUT status text,
+													  target_lsn pg_lsn,
+													  timeout int8 DEFAULT 0)
+  LANGUAGE internal AS 'pg_wal_replay_wait_status';
+
 CREATE OR REPLACE FUNCTION pg_logical_slot_get_changes(
     IN slot_name name, IN upto_lsn pg_lsn, IN upto_nchanges int, VARIADIC options text[] DEFAULT '{}',
     OUT lsn pg_lsn, OUT xid xid, OUT data text)
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 43f608d7a0a..36024a6e2c9 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -6647,6 +6647,17 @@
   proname => 'pg_wal_replay_wait', prokind => 'p', prorettype => 'void',
   proargtypes => 'pg_lsn int8', proargnames => '{target_lsn,timeout}',
   prosrc => 'pg_wal_replay_wait' },
+{ oid => '226',
+  descr => 'wait for the target LSN to be replayed on standby with an optional timeout without error throwing',
+  proname => 'pg_wal_replay_wait_no_error', prokind => 'p', prorettype => 'void',
+  proargtypes => 'pg_lsn int8', proargmodes => '{i,i}',
+  proargnames => '{target_lsn,timeout}',
+  prosrc => 'pg_wal_replay_wait_no_error' },
+{ oid => '388',
+  descr => 'return twait for the target LSN to be replayed on standby with an optional timeout and returning the waiting result',
+  proname => 'pg_wal_replay_wait_status', prorettype => 'text',
+  proargtypes => '',
+  prosrc => 'pg_wal_replay_wait_status' },
 
 { oid => '6224', descr => 'get resource managers loaded in system',
   proname => 'pg_get_wal_resource_managers', prorows => '50', proretset => 't',
diff --git a/src/test/recovery/t/043_wal_replay_wait.pl b/src/test/recovery/t/043_wal_replay_wait.pl
index 024f1fe6488..d5bf48895fa 100644
--- a/src/test/recovery/t/043_wal_replay_wait.pl
+++ b/src/test/recovery/t/043_wal_replay_wait.pl
@@ -77,6 +77,18 @@ $node_standby->psql(
 ok( $stderr =~ /timed out while waiting for target LSN/,
 	"get timeout on waiting for unreachable LSN");
 
+$output = $node_standby->safe_psql('postgres', qq[
+	CALL pg_wal_replay_wait_no_error('${lsn2}', 10);
+	SELECT pg_wal_replay_wait_status();]);
+ok( $output == "success",
+	"pg_wal_replay_wait_status() returns correct status after successful waiting");
+$output = $node_standby->psql(
+	'postgres',	qq[
+	CALL pg_wal_replay_wait_no_error('${lsn2}', 10);
+	SELECT pg_wal_replay_wait_status();]);
+ok( $output == "timeout",
+	"pg_wal_replay_wait_status() returns correct status after timeout");
+
 # 4. Check that pg_wal_replay_wait() triggers an error if called on primary,
 # within another function, or inside a transaction with an isolation level
 # higher than READ COMMITTED.
-- 
2.39.5 (Apple Git-154)