From 5962b11fb84358eb5df086a3059e64cda0a1b752 Mon Sep 17 00:00:00 2001
From: Alexander Korotkov <akorotkov@postgresql.org>
Date: Thu, 19 Sep 2024 15:34:18 +0300
Subject: [PATCH v6 4/4] Add 'no_error' argument to pg_wal_replay_wait()

This argument allow skipping throwing an error.  Instead, the result status
can be obtained using pg_wal_replay_wait_status() function.

Catversion is bumped.

Reported-by: Michael Paquier
Discussion: https://postgr.es/m/ZtUF17gF0pNpwZDI%40paquier.xyz
Reviewed-by: Pavel Borisov
---
 doc/src/sgml/func.sgml                     | 56 +++++++++++++++++++---
 src/backend/access/transam/xlogfuncs.c     | 38 +++++++++++++--
 src/backend/access/transam/xlogwait.c      |  3 +-
 src/backend/catalog/system_functions.sql   |  4 +-
 src/include/catalog/pg_proc.dat            |  7 ++-
 src/test/recovery/t/043_wal_replay_wait.pl | 22 +++++++++
 6 files changed, 117 insertions(+), 13 deletions(-)

diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
index ad663c94d77..1b68f284e5e 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -28989,12 +28989,15 @@ postgres=# SELECT '0/0'::pg_lsn + pd.segment_number * ps.setting::int + :offset
    </para>
 
    <table id="recovery-synchronization-procedure-table">
-    <title>Recovery Synchronization Procedure</title>
+    <title>Recovery Synchronization Procedure and Function</title>
     <tgroup cols="1">
      <thead>
       <row>
        <entry role="func_table_entry"><para role="func_signature">
-        Procedure
+        Procedure or Function
+       </para>
+       <para>
+        Type
        </para>
        <para>
         Description
@@ -29010,8 +29013,11 @@ postgres=# SELECT '0/0'::pg_lsn + pd.segment_number * ps.setting::int + :offset
         </indexterm>
         <function>pg_wal_replay_wait</function> (
           <parameter>target_lsn</parameter> <type>pg_lsn</type>,
-          <parameter>timeout</parameter> <type>bigint</type> <literal>DEFAULT</literal> <literal>0</literal>)
-        <returnvalue>void</returnvalue>
+          <parameter>timeout</parameter> <type>bigint</type> <literal>DEFAULT</literal> <literal>0</literal>,
+          <parameter>no_error</parameter> <type>bool</type> <literal>DEFAULT</literal> <literal>false</literal>)
+       </para>
+       <para>
+        Procedure
        </para>
        <para>
         Waits until recovery replays <literal>target_lsn</literal>.
@@ -29022,7 +29028,30 @@ postgres=# SELECT '0/0'::pg_lsn + pd.segment_number * ps.setting::int + :offset
         procedure waits until <literal>target_lsn</literal> is reached or
         the specified <parameter>timeout</parameter> has elapsed.
         On timeout, or if the server is promoted before
-        <literal>target_lsn</literal> is reached, an error is emitted.
+        <literal>target_lsn</literal> is reached, an error is emitted,
+        as soon as <parameter>no_error</parameter> is false.
+        If <parameter>no_error</parameter> is set to true, then the procedure
+        doesn't throw errors.  The last result status could be read
+        with <function>pg_wal_replay_wait_status</function>.
+       </para></entry>
+      </row>
+
+      <row>
+       <entry role="func_table_entry"><para role="func_signature">
+        <indexterm>
+         <primary>pg_wal_replay_wait_status</primary>
+        </indexterm>
+        <function>pg_wal_replay_wait_status</function> ()
+        <returnvalue>text</returnvalue>
+       </para>
+       <para>
+        Function
+       </para>
+       <para>
+        Returns the last result status for
+        <function>pg_wal_replay_wait</function> procedure.  The possible
+        values are <literal>success</literal>, <literal>timeout</literal>,
+        and <literal>not in recovery</literal>.
        </para></entry>
       </row>
      </tbody>
@@ -29044,7 +29073,8 @@ postgres=# SELECT '0/0'::pg_lsn + pd.segment_number * ps.setting::int + :offset
    <para>
     <function>pg_wal_replay_wait</function> should be called on standby.
     If a user calls <function>pg_wal_replay_wait</function> on primary, it
-    will error out.  However, if <function>pg_wal_replay_wait</function> is
+    will error out as soon as <parameter>no_error</parameter> is false.
+    However, if <function>pg_wal_replay_wait</function> is
     called on primary promoted from standby and <literal>target_lsn</literal>
     was already replayed, then <function>pg_wal_replay_wait</function> just
     exits immediately.
@@ -29090,6 +29120,20 @@ postgres=# CALL pg_wal_replay_wait('0/306EE20', 100);
 ERROR:  timed out while waiting for target LSN 0/306EE20 to be replayed; current replay LSN 0/306EA60
     </programlisting>
 
+   The same example uses <function>pg_wal_replay_wait</function> with
+   <parameter>no_error</parameter> set to true.  In this case, the result
+   status must be read with <function>pg_wal_replay_wait_status</function>.
+
+   <programlisting>
+postgres=# CALL pg_wal_replay_wait('0/306EE20', 100, true);
+CALL
+postgres=# SELECT pg_wal_replay_wait_status();
+ pg_wal_replay_wait_status
+---------------------------
+ timeout
+(1 row)
+    </programlisting>
+
    </para>
 
    <para>
diff --git a/src/backend/access/transam/xlogfuncs.c b/src/backend/access/transam/xlogfuncs.c
index ddca78d3717..bca1d395683 100644
--- a/src/backend/access/transam/xlogfuncs.c
+++ b/src/backend/access/transam/xlogfuncs.c
@@ -751,15 +751,18 @@ pg_promote(PG_FUNCTION_ARGS)
 	PG_RETURN_BOOL(false);
 }
 
+static WaitLSNResult lastWaitLSNResult = WAIT_LSN_RESULT_SUCCESS;
+
 /*
- * Waits until recovery replays the target LSN with optional timeout.
+ * Waits until recovery replays the target LSN with optional timeout.  Unless
+ * 'no_error' provided throws 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);
-	WaitLSNResult result;
+	bool		no_error = PG_GETARG_BOOL(2);
 
 	if (timeout < 0)
 		ereport(ERROR,
@@ -800,13 +803,16 @@ pg_wal_replay_wait(PG_FUNCTION_ARGS)
 	 */
 	Assert(MyProc->xmin == InvalidTransactionId);
 
-	result = WaitForLSNReplay(target_lsn, timeout);
+	lastWaitLSNResult = WaitForLSNReplay(target_lsn, timeout);
+
+	if (no_error)
+		PG_RETURN_VOID();
 
 	/*
 	 * Process the result of WaitForLSNReplay().  Throw appropriate error if
 	 * needed.
 	 */
-	switch (result)
+	switch (lastWaitLSNResult)
 	{
 		case WAIT_LSN_RESULT_SUCCESS:
 			/* Nothing to do on success */
@@ -832,3 +838,27 @@ pg_wal_replay_wait(PG_FUNCTION_ARGS)
 
 	PG_RETURN_VOID();
 }
+
+Datum
+pg_wal_replay_wait_status(PG_FUNCTION_ARGS)
+{
+	const char *result_string = "";
+
+	/* Process the result of WaitForLSNReplay(). */
+	switch (lastWaitLSNResult)
+	{
+		case WAIT_LSN_RESULT_SUCCESS:
+			result_string = "success";
+			break;
+
+		case WAIT_LSN_RESULT_TIMEOUT:
+			result_string = "timeout";
+			break;
+
+		case WAIT_LSN_RESULT_NOT_IN_RECOVERY:
+			result_string = "not in recovery";
+			break;
+	}
+
+	PG_RETURN_TEXT_P(cstring_to_text(result_string));
+}
diff --git a/src/backend/access/transam/xlogwait.c b/src/backend/access/transam/xlogwait.c
index 58fb10aa5a8..8860a9c73da 100644
--- a/src/backend/access/transam/xlogwait.c
+++ b/src/backend/access/transam/xlogwait.c
@@ -2,7 +2,8 @@
  *
  * xlogwait.c
  *	  Implements waiting for the given replay LSN, which is used in
- *	  CALL pg_wal_replay_wait(target_lsn pg_lsn, timeout float8).
+ *	  CALL pg_wal_replay_wait(target_lsn pg_lsn,
+ *							  timeout float8, no_error bool).
  *
  * Copyright (c) 2024, PostgreSQL Global Development Group
  *
diff --git a/src/backend/catalog/system_functions.sql b/src/backend/catalog/system_functions.sql
index 92b6aefe7aa..3fd1dfa945e 100644
--- a/src/backend/catalog/system_functions.sql
+++ b/src/backend/catalog/system_functions.sql
@@ -414,7 +414,9 @@ CREATE OR REPLACE FUNCTION
   json_populate_recordset(base anyelement, from_json json, use_json_as_text boolean DEFAULT false)
   RETURNS SETOF anyelement LANGUAGE internal STABLE ROWS 100  AS 'json_populate_recordset' PARALLEL SAFE;
 
-CREATE OR REPLACE PROCEDURE pg_wal_replay_wait(target_lsn pg_lsn, timeout int8 DEFAULT 0)
+CREATE OR REPLACE PROCEDURE pg_wal_replay_wait(target_lsn pg_lsn,
+											   timeout int8 DEFAULT 0,
+											   no_error bool DEFAULT false)
   LANGUAGE internal AS 'pg_wal_replay_wait';
 
 CREATE OR REPLACE FUNCTION pg_logical_slot_get_changes(
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 7c0b74fe055..4c8adab567f 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -6665,8 +6665,13 @@
 { oid => '8593',
   descr => 'wait for the target LSN to be replayed on standby with an optional timeout',
   proname => 'pg_wal_replay_wait', prokind => 'p', prorettype => 'void',
-  proargtypes => 'pg_lsn int8', proargnames => '{target_lsn,timeout}',
+  proargtypes => 'pg_lsn int8 bool', proargnames => '{target_lsn,timeout,no_error}',
   prosrc => 'pg_wal_replay_wait' },
+{ oid => '8594',
+  descr => 'the last result for pg_wal_replay_wait()',
+  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 cf77a9eec70..74769acd5e0 100644
--- a/src/test/recovery/t/043_wal_replay_wait.pl
+++ b/src/test/recovery/t/043_wal_replay_wait.pl
@@ -77,6 +77,20 @@ $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('${lsn2}', 10, true);
+	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('${lsn2}', 10, true);
+	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.
@@ -193,6 +207,14 @@ $node_standby->safe_psql('postgres', "CALL pg_wal_replay_wait('${lsn5}');");
 
 ok(1, 'wait for already replayed LSN exits immediately even after promotion');
 
+$output = $node_standby->psql(
+	'postgres', qq[
+	CALL pg_wal_replay_wait('${lsn4}', 10, true);
+	SELECT pg_wal_replay_wait_status();]);
+ok( $output == "not in recovery",
+	"pg_wal_replay_wait_status() returns correct status after standby promotion"
+);
+
 $node_standby->stop;
 $node_primary->stop;
 
-- 
2.39.5 (Apple Git-154)

