diff --git a/src/backend/access/transam/xlogfuncs.c b/src/backend/access/transam/xlogfuncs.c
index 4e46baaebd..3e3d2bb618 100644
--- a/src/backend/access/transam/xlogfuncs.c
+++ b/src/backend/access/transam/xlogfuncs.c
@@ -23,15 +23,18 @@
 #include "access/xlogbackup.h"
 #include "access/xlogrecovery.h"
 #include "catalog/pg_type.h"
+#include "commands/waitlsn.h"
 #include "funcapi.h"
 #include "miscadmin.h"
 #include "pgstat.h"
 #include "replication/walreceiver.h"
 #include "storage/fd.h"
+#include "storage/proc.h"
 #include "storage/standby.h"
 #include "utils/builtins.h"
 #include "utils/memutils.h"
 #include "utils/pg_lsn.h"
+#include "utils/snapmgr.h"
 #include "utils/timestamp.h"
 
 /*
@@ -747,3 +750,56 @@ pg_promote(PG_FUNCTION_ARGS)
 						   wait_seconds)));
 	PG_RETURN_BOOL(false);
 }
+
+/*
+ * Waits until recovery replays the target LSN with optional timeout.
+ */
+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")));
+
+	/*
+	 * 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.
+	 * Otherwise, our snapshot could prevent the replay of WAL records
+	 * 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
+	 * 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())
+		PopActiveSnapshot();
+
+	/*
+	 * At second, invalidate a catalog snapshot if any.  And we should be done
+	 * with the preparation.
+	 */
+	InvalidateCatalogSnapshot();
+
+	/* Give up if there is still an active or registered snapshot. */
+	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.")));
+
+	/*
+	 * As the result we should hold no snapshot, and correspondingly our xmin
+	 * should be unset.
+	 */
+	Assert(MyProc->xmin == InvalidTransactionId);
+
+	(void) WaitForLSNReplay(target_lsn, timeout);
+
+	PG_RETURN_VOID();
+}
diff --git a/src/backend/commands/waitlsn.c b/src/backend/commands/waitlsn.c
index ffa83a2000..501938f433 100644
--- a/src/backend/commands/waitlsn.c
+++ b/src/backend/commands/waitlsn.c
@@ -217,7 +217,7 @@ WaitLSNCleanup(void)
  * Wait using MyLatch till the given LSN is replayed, the postmaster dies or
  * timeout happens.
  */
-static void
+void
 WaitForLSNReplay(XLogRecPtr targetLSN, int64 timeout)
 {
 	XLogRecPtr	currentLSN;
@@ -336,53 +336,3 @@ WaitForLSNReplay(XLogRecPtr targetLSN, int64 timeout)
 						LSN_FORMAT_ARGS(currentLSN))));
 	}
 }
-
-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")));
-
-	/*
-	 * 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.
-	 * Otherwise, our snapshot could prevent the replay of WAL records
-	 * 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
-	 * 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())
-		PopActiveSnapshot();
-
-	/*
-	 * At second, invalidate a catalog snapshot if any.  And we should be done
-	 * with the preparation.
-	 */
-	InvalidateCatalogSnapshot();
-
-	/* Give up if there is still an active or registered snapshot. */
-	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.")));
-
-	/*
-	 * As the result we should hold no snapshot, and correspondingly our xmin
-	 * should be unset.
-	 */
-	Assert(MyProc->xmin == InvalidTransactionId);
-
-	(void) WaitForLSNReplay(target_lsn, timeout);
-
-	PG_RETURN_VOID();
-}
diff --git a/src/include/commands/waitlsn.h b/src/include/commands/waitlsn.h
index f719feadb0..bb5ac858dc 100644
--- a/src/include/commands/waitlsn.h
+++ b/src/include/commands/waitlsn.h
@@ -76,5 +76,6 @@ extern Size WaitLSNShmemSize(void);
 extern void WaitLSNShmemInit(void);
 extern void WaitLSNSetLatches(XLogRecPtr currentLSN);
 extern void WaitLSNCleanup(void);
+extern void WaitForLSNReplay(XLogRecPtr targetLSN, int64 timeout);
 
 #endif							/* WAIT_LSN_H */