From 4728be13bc17183f9869b1c040d5c72d2969e736 Mon Sep 17 00:00:00 2001
From: Amul Sul <amul.sul@enterprisedb.com>
Date: Tue, 14 Sep 2021 01:30:54 -0400
Subject: [PATCH v1] Deduplicate code updating ControleFile's DBState only.

---
 src/backend/access/transam/xlog.c | 53 +++++++++++++------------------
 src/include/access/xlog.h         |  3 ++
 2 files changed, 25 insertions(+), 31 deletions(-)

diff --git a/src/backend/access/transam/xlog.c b/src/backend/access/transam/xlog.c
index e51a7a749da..381a3576d5e 100644
--- a/src/backend/access/transam/xlog.c
+++ b/src/backend/access/transam/xlog.c
@@ -4951,6 +4951,19 @@ UpdateControlFile(void)
 	update_controlfile(DataDir, ControlFile, true);
 }
 
+/*
+ * Useful to set only ControlFile's database state.
+ */
+void
+SetControlFileDBState(DBState state)
+{
+	LWLockAcquire(ControlFileLock, LW_EXCLUSIVE);
+	ControlFile->state = state;
+	ControlFile->time = (pg_time_t) time(NULL);
+	UpdateControlFile();
+	LWLockRelease(ControlFileLock);
+}
+
 /*
  * Returns the unique system identifier from control file.
  */
@@ -8072,29 +8085,18 @@ StartupXLOG(void)
 	/*
 	 * All done with end-of-recovery actions.
 	 *
-	 * Now allow backends to write WAL and update the control file status in
-	 * consequence.  SharedRecoveryState, that controls if backends can write
-	 * WAL, is updated while holding ControlFileLock to prevent other backends
-	 * to look at an inconsistent state of the control file in shared memory.
-	 * There is still a small window during which backends can write WAL and
-	 * the control file is still referring to a system not in DB_IN_PRODUCTION
-	 * state while looking at the on-disk control file.
+	 * Now allow backends to write WAL and update the control file status and
+	 * SharedRecoveryState, that controls if backends can write WAL.
 	 *
-	 * Also, we use info_lck to update SharedRecoveryState to ensure that
-	 * there are no race conditions concerning visibility of other recent
-	 * updates to shared memory.
+	 * We use info_lck to update SharedRecoveryState to ensure that there are no
+	 * race conditions concerning visibility of other recent updates to shared
+	 * memory.
 	 */
-	LWLockAcquire(ControlFileLock, LW_EXCLUSIVE);
-	ControlFile->state = DB_IN_PRODUCTION;
-	ControlFile->time = (pg_time_t) time(NULL);
-
+	SetControlFileDBState(DB_IN_PRODUCTION);
 	SpinLockAcquire(&XLogCtl->info_lck);
 	XLogCtl->SharedRecoveryState = RECOVERY_STATE_DONE;
 	SpinLockRelease(&XLogCtl->info_lck);
 
-	UpdateControlFile();
-	LWLockRelease(ControlFileLock);
-
 	/*
 	 * If there were cascading standby servers connected to us, nudge any wal
 	 * sender processes to notice that we've been promoted.
@@ -8952,13 +8954,7 @@ CreateCheckPoint(int flags)
 	START_CRIT_SECTION();
 
 	if (shutdown)
-	{
-		LWLockAcquire(ControlFileLock, LW_EXCLUSIVE);
-		ControlFile->state = DB_SHUTDOWNING;
-		ControlFile->time = (pg_time_t) time(NULL);
-		UpdateControlFile();
-		LWLockRelease(ControlFileLock);
-	}
+		SetControlFileDBState(DB_SHUTDOWNING);
 
 	/*
 	 * Let smgr prepare for checkpoint; this has to happen before we determine
@@ -9507,13 +9503,8 @@ CreateRestartPoint(int flags)
 
 		UpdateMinRecoveryPoint(InvalidXLogRecPtr, true);
 		if (flags & CHECKPOINT_IS_SHUTDOWN)
-		{
-			LWLockAcquire(ControlFileLock, LW_EXCLUSIVE);
-			ControlFile->state = DB_SHUTDOWNED_IN_RECOVERY;
-			ControlFile->time = (pg_time_t) time(NULL);
-			UpdateControlFile();
-			LWLockRelease(ControlFileLock);
-		}
+			SetControlFileDBState(DB_SHUTDOWNED_IN_RECOVERY);
+
 		return false;
 	}
 
diff --git a/src/include/access/xlog.h b/src/include/access/xlog.h
index 0a8ede700de..3a485d86647 100644
--- a/src/include/access/xlog.h
+++ b/src/include/access/xlog.h
@@ -253,6 +253,8 @@ typedef enum WALAvailability
 	WALAVAIL_REMOVED			/* WAL segment has been removed */
 } WALAvailability;
 
+/* Avoid header inclusion */
+enum DBState;
 struct XLogRecData;
 
 extern XLogRecPtr XLogInsertRecord(struct XLogRecData *rdata,
@@ -291,6 +293,7 @@ extern TimestampTz GetLatestXTime(void);
 extern TimestampTz GetCurrentChunkReplayStartTime(void);
 
 extern void UpdateControlFile(void);
+extern void SetControlFileDBState(enum DBState state);
 extern uint64 GetSystemIdentifier(void);
 extern char *GetMockAuthenticationNonce(void);
 extern bool DataChecksumsEnabled(void);
-- 
2.18.0

