From 756d28ab2284e6eb8c057dae45280a38212cec48 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddyforpostgres@gmail.com>
Date: Tue, 24 May 2022 16:05:10 +0000
Subject: [PATCH v1] Switch to stream mode from archive occasionally

This patch enables standby to switch to stream mode i.e. get
WAL from primary even before it fails to receive from archive
location. Currently, if receive from archive location fails, it
switches back to stream mode and by then WAL receiver could have
come up. Since fetching WAL from archvie isn't cheap always,
switch to stream mode occasionally.

Right now, this switching is based on wal_retrieve_retry_interval
i.e. if the standby has received the WAL from archive for
wal_retrieve_retry_interval after the last failed attempt from
the stream, it switches back to stream mode to see if it can
receivefrom primary, if yes it's a good bet otherwise falls back
to archive.
---
 src/backend/access/transam/xlogrecovery.c | 53 +++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/src/backend/access/transam/xlogrecovery.c b/src/backend/access/transam/xlogrecovery.c
index 6eba626420..3746e24a18 100644
--- a/src/backend/access/transam/xlogrecovery.c
+++ b/src/backend/access/transam/xlogrecovery.c
@@ -241,6 +241,7 @@ static XLogSource readSource = XLOG_FROM_ANY;
  * walreceiver restart.  This is only valid in XLOG_FROM_STREAM state.
  */
 static XLogSource currentSource = XLOG_FROM_ANY;
+static XLogSource lastSource = XLOG_FROM_ANY;
 static bool lastSourceFailed = false;
 static bool pendingWalRcvRestart = false;
 
@@ -3065,6 +3066,7 @@ ReadRecord(XLogPrefetcher *xlogprefetcher, int emode,
 				 * so that we will check the archive next.
 				 */
 				lastSourceFailed = false;
+				lastSource = currentSource;
 				currentSource = XLOG_FROM_ANY;
 
 				continue;
@@ -3375,11 +3377,15 @@ WaitForWALToBecomeAvailable(XLogRecPtr RecPtr, bool randAccess,
 	 *-------
 	 */
 	if (!InArchiveRecovery)
+	{
+		lastSource = currentSource;
 		currentSource = XLOG_FROM_PG_WAL;
+	}
 	else if (currentSource == XLOG_FROM_ANY ||
 			 (!StandbyMode && currentSource == XLOG_FROM_STREAM))
 	{
 		lastSourceFailed = false;
+		lastSource = currentSource;
 		currentSource = XLOG_FROM_ARCHIVE;
 	}
 
@@ -3432,7 +3438,15 @@ WaitForWALToBecomeAvailable(XLogRecPtr RecPtr, bool randAccess,
 					 * Move to XLOG_FROM_STREAM state, and set to start a
 					 * walreceiver if necessary.
 					 */
+					lastSource = currentSource;
 					currentSource = XLOG_FROM_STREAM;
+					/*
+					 * XXX: we might have to see if the WAL receiver is already
+					 * running before even we just go ahead and start it
+					 * relying only on startWalReceiver flag. The WAL receiver
+					 * could've come up by then, if yes, there can be multiple
+					 * WAL receivers???
+					 */
 					startWalReceiver = true;
 					break;
 
@@ -3475,6 +3489,7 @@ WaitForWALToBecomeAvailable(XLogRecPtr RecPtr, bool randAccess,
 					{
 						if (rescanLatestTimeLine(replayTLI, replayLSN))
 						{
+							lastSource = currentSource;
 							currentSource = XLOG_FROM_ARCHIVE;
 							break;
 						}
@@ -3512,6 +3527,7 @@ WaitForWALToBecomeAvailable(XLogRecPtr RecPtr, bool randAccess,
 						HandleStartupProcInterrupts();
 					}
 					last_fail_time = now;
+					lastSource = currentSource;
 					currentSource = XLOG_FROM_ARCHIVE;
 					break;
 
@@ -3527,7 +3543,10 @@ WaitForWALToBecomeAvailable(XLogRecPtr RecPtr, bool randAccess,
 			 * from the archive first.
 			 */
 			if (InArchiveRecovery)
+			{
+				lastSource = currentSource;
 				currentSource = XLOG_FROM_ARCHIVE;
+			}
 		}
 
 		if (currentSource != oldSource)
@@ -3562,6 +3581,40 @@ WaitForWALToBecomeAvailable(XLogRecPtr RecPtr, bool randAccess,
 				if (randAccess)
 					curFileTLI = 0;
 
+				/*
+				 * Try to stream WAL from primary after a specified amount of
+				 * time fetching from archive. This is because fetching WAL
+				 * from archive isn't always cheaper and the primary could have
+				 * come up meanwhile.
+				 */
+				if (StandbyMode && lastSource == XLOG_FROM_STREAM)
+				{
+					TimestampTz now;
+
+					now = GetCurrentTimestamp();
+
+					if (TimestampDifferenceExceeds(last_fail_time, now,
+												   wal_retrieve_retry_interval))
+					{
+						elog(DEBUG2,
+							 "switching WAL source from %s to %s after \"wal_retrieve_retry_interval\" %d milliseconds",
+							 xlogSourceNames[currentSource],
+							 xlogSourceNames[lastSource],
+							 wal_retrieve_retry_interval);
+
+						currentSource =	lastSource;
+						last_fail_time = 0;
+
+						/*
+						 * Treat this as a failure to read from current source,
+						 * even though it is actually not, so that the state
+						 * machine moves to read it from XLOG_FROM_STREAM.
+						 */
+						lastSourceFailed = true;
+						break;
+					}
+				}
+
 				/*
 				 * Try to restore the file from archive, or read an existing
 				 * file from pg_wal.
-- 
2.25.1