From 6c1ce4710dbad27f230c68e4089d4faac8e5f385 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathandbossart@gmail.com>
Date: Fri, 23 Dec 2022 16:53:38 -0800
Subject: [PATCH v4 2/3] Refactor code for restoring files via shell.

Presently, restore_command uses a different code path than
archive_cleanup_command and recovery_end_command.  These code paths
are similar and can be easily combined.
---
 src/backend/access/transam/shell_restore.c | 108 +++++++++++----------
 1 file changed, 58 insertions(+), 50 deletions(-)

diff --git a/src/backend/access/transam/shell_restore.c b/src/backend/access/transam/shell_restore.c
index 3ddcabd969..073e709e06 100644
--- a/src/backend/access/transam/shell_restore.c
+++ b/src/backend/access/transam/shell_restore.c
@@ -22,17 +22,19 @@
 #include "storage/ipc.h"
 #include "utils/wait_event.h"
 
-static void ExecuteRecoveryCommand(const char *command,
+static char *BuildCleanupCommand(const char *command,
+								 const char *lastRestartPointFileName);
+static bool ExecuteRecoveryCommand(const char *command,
 								   const char *commandName, bool failOnSignal,
-								   uint32 wait_event_info,
-								   const char *lastRestartPointFileName);
+								   bool exitOnSigterm, uint32 wait_event_info,
+								   int fail_elevel);
 
 bool
 shell_restore(const char *file, const char *path,
 			  const char *lastRestartPointFileName)
 {
 	char	   *cmd;
-	int			rc;
+	bool		ret;
 
 	/* Build the restore command to execute */
 	cmd = BuildRestoreCommand(recoveryRestoreCommand, path, file,
@@ -40,19 +42,6 @@ shell_restore(const char *file, const char *path,
 	if (cmd == NULL)
 		elog(ERROR, "could not build restore command \"%s\"", cmd);
 
-	ereport(DEBUG3,
-			(errmsg_internal("executing restore command \"%s\"", cmd)));
-
-	/*
-	 * Copy xlog from archival storage to XLOGDIR
-	 */
-	fflush(NULL);
-	pgstat_report_wait_start(WAIT_EVENT_RESTORE_COMMAND);
-	rc = system(cmd);
-	pgstat_report_wait_end();
-
-	pfree(cmd);
-
 	/*
 	 * Remember, we rollforward UNTIL the restore fails so failure here is
 	 * just part of the process... that makes it difficult to determine
@@ -77,60 +66,52 @@ shell_restore(const char *file, const char *path,
 	 *
 	 * We treat hard shell errors such as "command not found" as fatal, too.
 	 */
-	if (wait_result_is_signal(rc, SIGTERM))
-		proc_exit(1);
-
-	ereport(wait_result_is_any_signal(rc, true) ? FATAL : DEBUG2,
-			(errmsg("could not restore file \"%s\" from archive: %s",
-					file, wait_result_to_str(rc))));
+	ret = ExecuteRecoveryCommand(cmd, "restore_command", true, true,
+								 WAIT_EVENT_RESTORE_COMMAND, DEBUG2);
+	pfree(cmd);
 
-	return (rc == 0);
+	return ret;
 }
 
 void
 shell_archive_cleanup(const char *lastRestartPointFileName)
 {
-	ExecuteRecoveryCommand(archiveCleanupCommand, "archive_cleanup_command",
-						   false, WAIT_EVENT_ARCHIVE_CLEANUP_COMMAND,
-						   lastRestartPointFileName);
+	char	   *cmd = BuildCleanupCommand(archiveCleanupCommand,
+										  lastRestartPointFileName);
+
+	(void) ExecuteRecoveryCommand(cmd, "archive_cleanup_command", false, false,
+								  WAIT_EVENT_ARCHIVE_CLEANUP_COMMAND, WARNING);
+	pfree(cmd);
 }
 
 void
 shell_recovery_end(const char *lastRestartPointFileName)
 {
-	ExecuteRecoveryCommand(recoveryEndCommand, "recovery_end_command", true,
-						   WAIT_EVENT_RECOVERY_END_COMMAND,
-						   lastRestartPointFileName);
+	char	   *cmd = BuildCleanupCommand(recoveryEndCommand,
+										  lastRestartPointFileName);
+
+	(void) ExecuteRecoveryCommand(cmd, "recovery_end_command", true, false,
+								  WAIT_EVENT_RECOVERY_END_COMMAND, WARNING);
+	pfree(cmd);
 }
 
 /*
- * Attempt to execute an external shell command during recovery.
- *
- * 'command' is the shell command to be executed, 'commandName' is a
- * human-readable name describing the command emitted in the logs. If
- * 'failOnSignal' is true and the command is killed by a signal, a FATAL
- * error is thrown. Otherwise a WARNING is emitted.
- *
- * This is currently used for recovery_end_command and archive_cleanup_command.
+ * Build a recovery_end_command or archive_cleanup_command.  The return value
+ * is palloc'd.
  */
-static void
-ExecuteRecoveryCommand(const char *command, const char *commandName,
-					   bool failOnSignal, uint32 wait_event_info,
-					   const char *lastRestartPointFileName)
+static char *
+BuildCleanupCommand(const char *command, const char *lastRestartPointFileName)
 {
-	char		xlogRecoveryCmd[MAXPGPATH];
+	char	   *ret = (char *) palloc(MAXPGPATH);
 	char	   *dp;
 	char	   *endp;
 	const char *sp;
-	int			rc;
-
-	Assert(command && commandName);
 
 	/*
 	 * construct the command to be executed
 	 */
-	dp = xlogRecoveryCmd;
-	endp = xlogRecoveryCmd + MAXPGPATH - 1;
+	dp = ret;
+	endp = ret + MAXPGPATH - 1;
 	*endp = '\0';
 
 	for (sp = command; *sp; sp++)
@@ -166,6 +147,28 @@ ExecuteRecoveryCommand(const char *command, const char *commandName,
 	}
 	*dp = '\0';
 
+	return ret;
+}
+
+/*
+ * Attempt to execute an external shell command during recovery.
+ *
+ * 'command' is the shell command to be executed, 'commandName' is a
+ * human-readable name describing the command emitted in the logs. If
+ * 'failOnSignal' is true and the command is killed by a signal, a FATAL error
+ * is thrown. Otherwise, 'fail_elevel' is used for the log message.  If
+ * 'exitOnSigterm' is true and the command is killed by SIGTERM, we exit
+ * immediately.
+ *
+ * Returns whether the command succeeded.
+ */
+static bool
+ExecuteRecoveryCommand(const char *command, const char *commandName,
+					   bool failOnSignal, bool exitOnSigterm,
+					   uint32 wait_event_info, int fail_elevel)
+{
+	int		rc;
+
 	ereport(DEBUG3,
 			(errmsg_internal("executing %s \"%s\"", commandName, command)));
 
@@ -174,16 +177,19 @@ ExecuteRecoveryCommand(const char *command, const char *commandName,
 	 */
 	fflush(NULL);
 	pgstat_report_wait_start(wait_event_info);
-	rc = system(xlogRecoveryCmd);
+	rc = system(command);
 	pgstat_report_wait_end();
 
 	if (rc != 0)
 	{
+		if (exitOnSigterm && wait_result_is_signal(rc, SIGTERM))
+			proc_exit(1);
+
 		/*
 		 * If the failure was due to any sort of signal, it's best to punt and
 		 * abort recovery.  See comments in shell_restore().
 		 */
-		ereport((failOnSignal && wait_result_is_any_signal(rc, true)) ? FATAL : WARNING,
+		ereport((failOnSignal && wait_result_is_any_signal(rc, true)) ? FATAL : fail_elevel,
 		/*------
 		   translator: First %s represents a postgresql.conf parameter name like
 		  "recovery_end_command", the 2nd is the value of that parameter, the
@@ -191,4 +197,6 @@ ExecuteRecoveryCommand(const char *command, const char *commandName,
 				(errmsg("%s \"%s\": %s", commandName,
 						command, wait_result_to_str(rc))));
 	}
+
+	return (rc == 0);
 }
-- 
2.25.1

