From 4aa18feb38ec7a77ebc8215ea4207071ce81d9d5 Mon Sep 17 00:00:00 2001
From: Amul Sul <amul.sul@enterprisedb.com>
Date: Mon, 20 Sep 2021 01:50:59 -0400
Subject: [PATCH v2] Deduplicate code updating ControleFile's DBState only.

---
 src/backend/access/transam/xlog.c | 48 ++++++++++++++++++++-----------
 1 file changed, 31 insertions(+), 17 deletions(-)

diff --git a/src/backend/access/transam/xlog.c b/src/backend/access/transam/xlog.c
index e51a7a749da..df1202e07da 100644
--- a/src/backend/access/transam/xlog.c
+++ b/src/backend/access/transam/xlog.c
@@ -931,6 +931,8 @@ static bool rescanLatestTimeLine(void);
 static void InitControlFile(uint64 sysidentifier);
 static void WriteControlFile(void);
 static void ReadControlFile(void);
+static void SetControlFileDBStateWithLock(DBState state);
+static void SetControlFileDBState(DBState state);
 static char *str_time(pg_time_t tnow);
 static void SetPromoteIsTriggered(void);
 static bool CheckForStandbyTrigger(void);
@@ -4951,6 +4953,31 @@ UpdateControlFile(void)
 	update_controlfile(DataDir, ControlFile, true);
 }
 
+/*
+ * Useful to set only ControlFile's database state when a caller has already
+ * acquired exclusive ControlFileLock.
+ */
+static void
+SetControlFileDBStateWithLock(DBState state)
+{
+	Assert(LWLockHeldByMeInMode(ControlFileLock, LW_EXCLUSIVE));
+
+	ControlFile->state = state;
+	ControlFile->time = (pg_time_t) time(NULL);
+	UpdateControlFile();
+}
+
+/*
+ * Useful to set only ControlFile's database state.
+ */
+static void
+SetControlFileDBState(DBState state)
+{
+	LWLockAcquire(ControlFileLock, LW_EXCLUSIVE);
+	SetControlFileDBStateWithLock(state);
+	LWLockRelease(ControlFileLock);
+}
+
 /*
  * Returns the unique system identifier from control file.
  */
@@ -8085,14 +8112,12 @@ StartupXLOG(void)
 	 * updates to shared memory.
 	 */
 	LWLockAcquire(ControlFileLock, LW_EXCLUSIVE);
-	ControlFile->state = DB_IN_PRODUCTION;
-	ControlFile->time = (pg_time_t) time(NULL);
+	SetControlFileDBStateWithLock(DB_IN_PRODUCTION);
 
 	SpinLockAcquire(&XLogCtl->info_lck);
 	XLogCtl->SharedRecoveryState = RECOVERY_STATE_DONE;
 	SpinLockRelease(&XLogCtl->info_lck);
 
-	UpdateControlFile();
 	LWLockRelease(ControlFileLock);
 
 	/*
@@ -8952,13 +8977,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 +9526,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;
 	}
 
-- 
2.18.0