From 6ebdff0cf9b5314fcee8227bcab92e21f8d8a866 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddyforpostgres@gmail.com>
Date: Thu, 7 Apr 2022 13:23:04 +0000
Subject: [PATCH v1] Allow pg_walfile_{name, name_offset} to run in recovery

Right now, pg_walfile_name and pg_walfile_name_offset don't run
while the server is in recovery (standby, PITR, crash). This
reduces their usability if the server opens up for read-only
connections in recovery.

This patch enables them to compute WAL file name even in recovery
with a timeline ID of last received and flushed WAL record of WAL
receiver if it's streaming, otherwise of the last replayed WAL
record.
---
 src/backend/access/transam/xlogfuncs.c     | 68 +++++++++++++++++-----
 src/backend/replication/walreceiverfuncs.c | 15 +++++
 src/include/replication/walreceiver.h      |  1 +
 3 files changed, 68 insertions(+), 16 deletions(-)

diff --git a/src/backend/access/transam/xlogfuncs.c b/src/backend/access/transam/xlogfuncs.c
index b61ae6c0b4..65c04caec3 100644
--- a/src/backend/access/transam/xlogfuncs.c
+++ b/src/backend/access/transam/xlogfuncs.c
@@ -44,6 +44,8 @@
 static StringInfo label_file;
 static StringInfo tblspc_map_file;
 
+static TimeLineID GetTLIForWALFileNameComputation(void);
+
 /*
  * pg_backup_start: set up for taking an on-line backup dump
  *
@@ -316,6 +318,50 @@ pg_last_wal_replay_lsn(PG_FUNCTION_ARGS)
 	PG_RETURN_LSN(recptr);
 }
 
+/*
+ * Get timeLine ID for computing WAL file name.
+ *
+ * When not in recovery, it returns the current timeline of the server.
+ *
+ * When in recovery and WAL receiver is streaming, it returns the timeline of
+ * the last WAL record received and flushed to disk.
+ *
+ * When in recovery and WAL receiver is not streaming, it returns the timeline
+ * of last WAL record successfully replayed.
+ */
+static TimeLineID
+GetTLIForWALFileNameComputation(void)
+{
+	TimeLineID	tli;
+
+	if (!RecoveryInProgress())
+		tli = GetWALInsertionTimeLine();
+	else
+	{
+		/*
+		 * When WAL receiver isn't streaming, do not read timeline ID from its
+		 * shared memory as we don't know yet whether it will come up and start
+		 * streaming.
+		 */
+		if (WalRcvGetState() == WALRCV_STREAMING)
+			(void) GetWalRcvFlushRecPtr(NULL, &tli);
+		else
+			(void) GetXLogReplayRecPtr(&tli);
+	}
+
+	/*
+	 * The computed timeline ID can be invalid, especially when in recovery.
+	 * Let's be cautious and emit an error, otherwise the caller may give wrong
+	 * WAL file name.
+	 */
+	if (tli == 0)
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("TimeLine ID %u is invalid", tli)));
+
+	return tli;
+}
+
 /*
  * Compute an xlog file name and decimal byte offset given a WAL location,
  * such as is returned by pg_backup_stop() or pg_switch_wal().
@@ -336,13 +382,9 @@ pg_walfile_name_offset(PG_FUNCTION_ARGS)
 	TupleDesc	resultTupleDesc;
 	HeapTuple	resultHeapTuple;
 	Datum		result;
+	TimeLineID	tli;
 
-	if (RecoveryInProgress())
-		ereport(ERROR,
-				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
-				 errmsg("recovery is in progress"),
-				 errhint("%s cannot be executed during recovery.",
-						 "pg_walfile_name_offset()")));
+	tli = GetTLIForWALFileNameComputation();
 
 	/*
 	 * Construct a tuple descriptor for the result row.  This must match this
@@ -360,8 +402,7 @@ pg_walfile_name_offset(PG_FUNCTION_ARGS)
 	 * xlogfilename
 	 */
 	XLByteToPrevSeg(locationpoint, xlogsegno, wal_segment_size);
-	XLogFileName(xlogfilename, GetWALInsertionTimeLine(), xlogsegno,
-				 wal_segment_size);
+	XLogFileName(xlogfilename, tli, xlogsegno, wal_segment_size);
 
 	values[0] = CStringGetTextDatum(xlogfilename);
 	isnull[0] = false;
@@ -394,17 +435,12 @@ pg_walfile_name(PG_FUNCTION_ARGS)
 	XLogSegNo	xlogsegno;
 	XLogRecPtr	locationpoint = PG_GETARG_LSN(0);
 	char		xlogfilename[MAXFNAMELEN];
+	TimeLineID	tli;
 
-	if (RecoveryInProgress())
-		ereport(ERROR,
-				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
-				 errmsg("recovery is in progress"),
-				 errhint("%s cannot be executed during recovery.",
-						 "pg_walfile_name()")));
+	tli = GetTLIForWALFileNameComputation();
 
 	XLByteToPrevSeg(locationpoint, xlogsegno, wal_segment_size);
-	XLogFileName(xlogfilename, GetWALInsertionTimeLine(), xlogsegno,
-				 wal_segment_size);
+	XLogFileName(xlogfilename, tli, xlogsegno, wal_segment_size);
 
 	PG_RETURN_TEXT_P(cstring_to_text(xlogfilename));
 }
diff --git a/src/backend/replication/walreceiverfuncs.c b/src/backend/replication/walreceiverfuncs.c
index 90798b9d53..0a17a5b4cd 100644
--- a/src/backend/replication/walreceiverfuncs.c
+++ b/src/backend/replication/walreceiverfuncs.c
@@ -171,6 +171,21 @@ WalRcvStreaming(void)
 		return false;
 }
 
+/*
+ * Get walreceiver state.
+ */
+WalRcvState
+WalRcvGetState(void)
+{
+	WalRcvState state;
+
+	SpinLockAcquire(&WalRcv->mutex);
+	state = WalRcv->walRcvState;
+	SpinLockRelease(&WalRcv->mutex);
+
+	return state;
+}
+
 /*
  * Stop walreceiver (if running) and wait for it to die.
  * Executed by the Startup process.
diff --git a/src/include/replication/walreceiver.h b/src/include/replication/walreceiver.h
index 92f73a55b8..d950425ab2 100644
--- a/src/include/replication/walreceiver.h
+++ b/src/include/replication/walreceiver.h
@@ -460,6 +460,7 @@ extern void WalRcvShmemInit(void);
 extern void ShutdownWalRcv(void);
 extern bool WalRcvStreaming(void);
 extern bool WalRcvRunning(void);
+extern WalRcvState WalRcvGetState(void);
 extern void RequestXLogStreaming(TimeLineID tli, XLogRecPtr recptr,
 								 const char *conninfo, const char *slotname,
 								 bool create_temp_slot);
-- 
2.25.1

