From 4e43cda2281d53c6fd4702a9d80553f0f9bd7458 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathandbossart@gmail.com>
Date: Fri, 3 Feb 2023 10:34:05 -0800
Subject: [PATCH v4 3/3] handle shutdown requests before and after restoring a
 file

---
 src/backend/access/transam/xlogarchive.c |  8 ++++++++
 src/backend/postmaster/startup.c         | 16 ++++++++++++----
 src/include/postmaster/startup.h         |  1 +
 3 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/src/backend/access/transam/xlogarchive.c b/src/backend/access/transam/xlogarchive.c
index 66312c816b..18cf5b032b 100644
--- a/src/backend/access/transam/xlogarchive.c
+++ b/src/backend/access/transam/xlogarchive.c
@@ -147,11 +147,19 @@ RestoreArchivedFile(char *path, const char *xlogfname,
 	else
 		XLogFileName(lastRestartPointFname, 0, 0L, wal_segment_size);
 
+	/*
+	 * Check for pending shutdown requests before and after restoring the file,
+	 * and exit if there is one.
+	 */
+	HandleStartupProcShutdownRequest();
+
 	/*
 	 * Copy xlog from archival storage to XLOGDIR
 	 */
 	ret = shell_restore(xlogfname, xlogpath, lastRestartPointFname);
 
+	HandleStartupProcShutdownRequest();
+
 	if (ret)
 	{
 		/*
diff --git a/src/backend/postmaster/startup.c b/src/backend/postmaster/startup.c
index 503eb1a5a6..499dc72404 100644
--- a/src/backend/postmaster/startup.c
+++ b/src/backend/postmaster/startup.c
@@ -195,8 +195,7 @@ HandleStartupProcInterrupts(void)
 	/*
 	 * Check if we were requested to exit without finishing recovery.
 	 */
-	if (shutdown_requested)
-		proc_exit(1);
+	HandleStartupProcShutdownRequest();
 
 	/*
 	 * Emergency bailout if postmaster has died.  This is to avoid the
@@ -220,6 +219,16 @@ HandleStartupProcInterrupts(void)
 		ProcessLogMemoryContextInterrupt();
 }
 
+/*
+ * Exit if there is a pending shutdown request.
+ */
+void
+HandleStartupProcShutdownRequest(void)
+{
+	if (shutdown_requested)
+		proc_exit(1);
+}
+
 
 /* --------------------------------
  *		signal handler routines
@@ -295,8 +304,7 @@ PreRestoreCommand(void)
 	 * shutdown request received just before this.
 	 */
 	in_restore_command = true;
-	if (shutdown_requested)
-		proc_exit(1);
+	HandleStartupProcShutdownRequest();
 }
 
 void
diff --git a/src/include/postmaster/startup.h b/src/include/postmaster/startup.h
index dd957f9291..8eea0490ff 100644
--- a/src/include/postmaster/startup.h
+++ b/src/include/postmaster/startup.h
@@ -29,6 +29,7 @@ extern void HandleStartupProcInterrupts(void);
 extern void StartupProcessMain(void) pg_attribute_noreturn();
 extern void PreRestoreCommand(void);
 extern void PostRestoreCommand(void);
+extern void HandleStartupProcShutdownRequest(void);
 extern bool IsPromoteSignaled(void);
 extern void ResetPromoteSignaled(void);
 
-- 
2.25.1