From: Nathan Bossart <nathandbossart@gmail.com>
Date: Thu, 2 Feb 2023 12:04:12 -0800
Subject: [PATCH v2 1/1] stopgap fix for restore_command

---
 src/backend/access/transam/shell_restore.c | 15 +++++++++++-
 src/backend/access/transam/xlogarchive.c   |  7 ------
 src/backend/postmaster/startup.c           | 27 +++++++++++++++++-----
 src/include/postmaster/startup.h           |  3 +--
 4 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/src/backend/access/transam/shell_restore.c b/src/backend/access/transam/shell_restore.c
index 8458209f49..8fc3e86a10 100644
--- a/src/backend/access/transam/shell_restore.c
+++ b/src/backend/access/transam/shell_restore.c
@@ -21,6 +21,8 @@
 #include "access/xlogarchive.h"
 #include "access/xlogrecovery.h"
 #include "common/percentrepl.h"
+#include "miscadmin.h"
+#include "postmaster/startup.h"
 #include "storage/ipc.h"
 #include "utils/wait_event.h"
 
@@ -146,7 +148,18 @@ ExecuteRecoveryCommand(const char *command, const char *commandName,
 	 */
 	fflush(NULL);
 	pgstat_report_wait_start(wait_event_info);
-	rc = system(command);
+
+	/*
+	 * When exitOnSigterm is set and we are in the startup process, use the
+	 * special wrapper for system() that enables exiting immediately upon
+	 * receiving SIGTERM.  This ensures we can break out of system() if
+	 * required.
+	 */
+	if (exitOnSigterm && MyBackendType == B_STARTUP)
+		rc = RunInterruptibleShellCommand(command);
+	else
+		rc = system(command);
+
 	pgstat_report_wait_end();
 
 	if (rc != 0)
diff --git a/src/backend/access/transam/xlogarchive.c b/src/backend/access/transam/xlogarchive.c
index 4b89addf97..66312c816b 100644
--- a/src/backend/access/transam/xlogarchive.c
+++ b/src/backend/access/transam/xlogarchive.c
@@ -147,18 +147,11 @@ RestoreArchivedFile(char *path, const char *xlogfname,
 	else
 		XLogFileName(lastRestartPointFname, 0, 0L, wal_segment_size);
 
-	/*
-	 * Check signals before restore command and reset afterwards.
-	 */
-	PreRestoreCommand();
-
 	/*
 	 * Copy xlog from archival storage to XLOGDIR
 	 */
 	ret = shell_restore(xlogfname, xlogpath, lastRestartPointFname);
 
-	PostRestoreCommand();
-
 	if (ret)
 	{
 		/*
diff --git a/src/backend/postmaster/startup.c b/src/backend/postmaster/startup.c
index 8786186898..aa94430c6f 100644
--- a/src/backend/postmaster/startup.c
+++ b/src/backend/postmaster/startup.c
@@ -273,9 +273,24 @@ StartupProcessMain(void)
 	proc_exit(0);
 }
 
-void
-PreRestoreCommand(void)
+/*
+ * This is a wrapper for system() that enables exiting immediately on SIGTERM.
+ * It is intended for use with restore_command since there isn't a good way to
+ * break out while it is executing.  Note that this behavior only works in the
+ * startup process.
+ *
+ * NB: Since we might call proc_exit() in a signal handler here, it is
+ * imperative that that nothing but the system() call happens between setting
+ * and resetting in_restore_command.  Any additional code must go before or
+ * after this section.
+ */
+int
+RunInterruptibleShellCommand(const char *command)
 {
+	int ret;
+
+	Assert(MyBackendType == B_STARTUP);
+
 	/*
 	 * Set in_restore_command to tell the signal handler that we should exit
 	 * right away on SIGTERM. We know that we're at a safe point to do that.
@@ -285,12 +300,12 @@ PreRestoreCommand(void)
 	in_restore_command = true;
 	if (shutdown_requested)
 		proc_exit(1);
-}
 
-void
-PostRestoreCommand(void)
-{
+	ret = system(command);
+
 	in_restore_command = false;
+
+	return ret;
 }
 
 bool
diff --git a/src/include/postmaster/startup.h b/src/include/postmaster/startup.h
index dd957f9291..5188f49d21 100644
--- a/src/include/postmaster/startup.h
+++ b/src/include/postmaster/startup.h
@@ -27,8 +27,7 @@ extern PGDLLIMPORT int log_startup_progress_interval;
 
 extern void HandleStartupProcInterrupts(void);
 extern void StartupProcessMain(void) pg_attribute_noreturn();
-extern void PreRestoreCommand(void);
-extern void PostRestoreCommand(void);
+extern int RunInterruptibleShellCommand(const char *command);
 extern bool IsPromoteSignaled(void);
 extern void ResetPromoteSignaled(void);
 
-- 
2.25.1


--LQksG6bCIzRHxTLp--