From a05c3343e0bd6fe339c944f6b0cde64ceb46a0b3 Mon Sep 17 00:00:00 2001
From: Alexey Kondratov <kondratov.aleksey@gmail.com>
Date: Tue, 19 Feb 2019 19:14:53 +0300
Subject: [PATCH v11] pg_rewind: options to use restore_command from command
 line or cluster config

Previously, when pg_rewind could not find required WAL files in the
target data directory the rewind process would fail. One had to
manually figure out which of required WAL files have already moved to
the archival storage and copy them back.

This patch adds possibility to specify restore_command via command
line option or use one specified inside postgresql.conf. Specified
restore_command will be used for automatic retrieval of missing WAL
files from archival storage.
---
 doc/src/sgml/ref/pg_rewind.sgml       |  49 +++++++-
 src/bin/pg_rewind/parsexlog.c         | 164 +++++++++++++++++++++++++-
 src/bin/pg_rewind/pg_rewind.c         | 118 +++++++++++++++---
 src/bin/pg_rewind/pg_rewind.h         |   6 +-
 src/bin/pg_rewind/t/001_basic.pl      |   4 +-
 src/bin/pg_rewind/t/002_databases.pl  |   4 +-
 src/bin/pg_rewind/t/003_extrafiles.pl |   4 +-
 src/bin/pg_rewind/t/RewindTest.pm     | 105 ++++++++++++++++-
 8 files changed, 416 insertions(+), 38 deletions(-)

diff --git a/doc/src/sgml/ref/pg_rewind.sgml b/doc/src/sgml/ref/pg_rewind.sgml
index 42d29edd4e..b601a5c7e4 100644
--- a/doc/src/sgml/ref/pg_rewind.sgml
+++ b/doc/src/sgml/ref/pg_rewind.sgml
@@ -66,11 +66,12 @@ PostgreSQL documentation
    can be found either on the target timeline, the source timeline, or their common
    ancestor. In the typical failover scenario where the target cluster was
    shut down soon after the divergence, this is not a problem, but if the
-   target cluster ran for a long time after the divergence, the old WAL
-   files might no longer be present. In that case, they can be manually
-   copied from the WAL archive to the <filename>pg_wal</filename> directory, or
-   fetched on startup by configuring <xref linkend="guc-primary-conninfo"/> or
-   <xref linkend="guc-restore-command"/>.  The use of
+   target cluster ran for a long time after the divergence, its old WAL
+   files might no longer be present. In this case, you can manually copy them
+   from the WAL archive to the <filename>pg_wal</filename> directory, or run
+   <application>pg_rewind</application> with the <literal>-c</literal> or
+   <literal>-C</literal> option to automatically retrieve them from the WAL
+   archive. The use of
    <application>pg_rewind</application> is not limited to failover, e.g.  a standby
    server can be promoted, run some write transactions, and then rewinded
    to become a standby again.
@@ -232,6 +233,39 @@ PostgreSQL documentation
       </listitem>
      </varlistentry>
 
+     <varlistentry>
+      <term><option>-c</option></term>
+      <term><option>--restore-target-wal</option></term>
+      <listitem>
+       <para>
+        Use the <varname>restore_command</varname> defined in
+        <filename>postgresql.conf</filename> to retrieve WAL files from
+        the WAL archive if these files are no longer available in the
+        <filename>pg_wal</filename> directory of the target cluster.
+       </para>
+       <para>
+        This option cannot be used together with <option>--target-restore-command</option>.
+       </para>
+      </listitem>
+     </varlistentry>
+
+     <varlistentry>
+      <term><option>-C <replaceable class="parameter">restore_command</replaceable></option></term>
+      <term><option>--target-restore-command=<replaceable class="parameter">restore_command</replaceable></option></term>
+      <listitem>
+       <para>
+        Specifies the <varname>restore_command</varname> to use for retrieving
+        WAL files from the WAL archive if these files are no longer available
+        in the <filename>pg_wal</filename> directory of the target cluster.
+       </para>
+       <para>
+        If <varname>restore_command</varname> is already set in
+        <filename>postgresql.conf</filename>, you can provide the
+        <option>--restore-target-wal</option> option instead.
+       </para>
+      </listitem>
+     </varlistentry>
+
      <varlistentry>
       <term><option>--debug</option></term>
       <listitem>
@@ -318,7 +352,10 @@ GRANT EXECUTE ON function pg_catalog.pg_read_binary_file(text, bigint, bigint, b
       history forked off from the target cluster. For each WAL record,
       record each data block that was touched. This yields a list of all
       the data blocks that were changed in the target cluster, after the
-      source cluster forked off.
+      source cluster forked off. If some of the WAL files are no longer
+      available, try re-running <application>pg_rewind</application> with
+      the <option>-c</option> or <option>-C</option> option to search
+      for the missing files in the WAL archive.
      </para>
     </step>
     <step>
diff --git a/src/bin/pg_rewind/parsexlog.c b/src/bin/pg_rewind/parsexlog.c
index 1d03375a3e..fd67790ac3 100644
--- a/src/bin/pg_rewind/parsexlog.c
+++ b/src/bin/pg_rewind/parsexlog.c
@@ -12,6 +12,7 @@
 #include "postgres_fe.h"
 
 #include <unistd.h>
+#include <sys/stat.h>
 
 #include "access/rmgr.h"
 #include "access/xlog_internal.h"
@@ -41,6 +42,7 @@ static char xlogfpath[MAXPGPATH];
 
 typedef struct XLogPageReadPrivate
 {
+	const char *restoreCommand;
 	int			tliIndex;
 } XLogPageReadPrivate;
 
@@ -48,6 +50,9 @@ static int	SimpleXLogPageRead(XLogReaderState *xlogreader,
 							   XLogRecPtr targetPagePtr,
 							   int reqLen, XLogRecPtr targetRecPtr, char *readBuf);
 
+static int RestoreArchivedWAL(const char *path, const char *xlogfname,
+				   off_t expectedSize, const char *restoreCommand);
+
 /*
  * Read WAL from the datadir/pg_wal, starting from 'startpoint' on timeline
  * index 'tliIndex' in target timeline history, until 'endpoint'. Make note of
@@ -55,7 +60,7 @@ static int	SimpleXLogPageRead(XLogReaderState *xlogreader,
  */
 void
 extractPageMap(const char *datadir, XLogRecPtr startpoint, int tliIndex,
-			   XLogRecPtr endpoint)
+			   XLogRecPtr endpoint, const char *restore_command)
 {
 	XLogRecord *record;
 	XLogReaderState *xlogreader;
@@ -63,6 +68,7 @@ extractPageMap(const char *datadir, XLogRecPtr startpoint, int tliIndex,
 	XLogPageReadPrivate private;
 
 	private.tliIndex = tliIndex;
+	private.restoreCommand = restore_command;
 	xlogreader = XLogReaderAllocate(WalSegSz, datadir, &SimpleXLogPageRead,
 									&private);
 	if (xlogreader == NULL)
@@ -148,7 +154,7 @@ readOneRecord(const char *datadir, XLogRecPtr ptr, int tliIndex)
 void
 findLastCheckpoint(const char *datadir, XLogRecPtr forkptr, int tliIndex,
 				   XLogRecPtr *lastchkptrec, TimeLineID *lastchkpttli,
-				   XLogRecPtr *lastchkptredo)
+				   XLogRecPtr *lastchkptredo, const char *restoreCommand)
 {
 	/* Walk backwards, starting from the given record */
 	XLogRecord *record;
@@ -172,6 +178,7 @@ findLastCheckpoint(const char *datadir, XLogRecPtr forkptr, int tliIndex,
 	}
 
 	private.tliIndex = tliIndex;
+	private.restoreCommand = restoreCommand;
 	xlogreader = XLogReaderAllocate(WalSegSz, datadir, &SimpleXLogPageRead,
 									&private);
 	if (xlogreader == NULL)
@@ -282,8 +289,29 @@ SimpleXLogPageRead(XLogReaderState *xlogreader, XLogRecPtr targetPagePtr,
 
 		if (xlogreadfd < 0)
 		{
-			pg_log_error("could not open file \"%s\": %m", xlogfpath);
-			return -1;
+			/*
+			 * If we have no restore_command to execute, then exit.
+			 */
+			if (private->restoreCommand == NULL)
+			{
+				pg_log_error("could not open file \"%s\": %m", xlogfpath);
+				return -1;
+			}
+
+			/*
+			 * Since we have restore_command to execute, then try to retrieve
+			 * missing WAL file from the archive.
+			 */
+			xlogreadfd = RestoreArchivedWAL(xlogreader->segcxt.ws_dir,
+											xlogfname,
+											WalSegSz,
+											private->restoreCommand);
+
+			if (xlogreadfd < 0)
+				return -1;
+			else
+				pg_log_debug("using file \"%s\" restored from archive",
+							 xlogfpath);
 		}
 	}
 
@@ -397,3 +425,131 @@ extractPageInfo(XLogReaderState *record)
 		process_block_change(forknum, rnode, blkno);
 	}
 }
+
+/*
+ * Attempt to retrieve the specified file from off-line archival storage.
+ * If successful return a file descriptor of restored WAL file, else
+ * return -1.
+ *
+ * For fixed-size files, the caller may pass the expected size as an
+ * additional crosscheck on successful recovery. If the file size is not
+ * known, set expectedSize = 0.
+ */
+static int
+RestoreArchivedWAL(const char *path, const char *xlogfname,
+				   off_t expectedSize, const char *restoreCommand)
+{
+	char		xlogpath[MAXPGPATH],
+				xlogRestoreCmd[MAXPGPATH],
+			   *dp,
+			   *endp;
+	const char *sp;
+	int			rc,
+				xlogfd;
+	struct stat stat_buf;
+
+	snprintf(xlogpath, MAXPGPATH, "%s/" XLOGDIR "/%s", path, xlogfname);
+
+	/*
+	 * Construct the command to be executed.
+	 */
+	dp = xlogRestoreCmd;
+	endp = xlogRestoreCmd + MAXPGPATH - 1;
+	*endp = '\0';
+
+	for (sp = restoreCommand; *sp; sp++)
+	{
+		if (*sp == '%')
+		{
+			switch (sp[1])
+			{
+				case 'p':
+					/* %p: relative path of target file */
+					sp++;
+					StrNCpy(dp, xlogpath, endp - dp);
+					make_native_path(dp);
+					dp += strlen(dp);
+					break;
+				case 'f':
+					/* %f: filename of desired file */
+					sp++;
+					StrNCpy(dp, xlogfname, endp - dp);
+					dp += strlen(dp);
+					break;
+				case 'r':
+					/* %r: filename of last restartpoint */
+					pg_fatal("restore_command with %%r cannot be used with pg_rewind.");
+					break;
+				case '%':
+					/* convert %% to a single % */
+					sp++;
+					if (dp < endp)
+						*dp++ = *sp;
+					break;
+				default:
+					/* otherwise treat the % as not special */
+					if (dp < endp)
+						*dp++ = *sp;
+					break;
+			}
+		}
+		else
+		{
+			if (dp < endp)
+				*dp++ = *sp;
+		}
+	}
+	*dp = '\0';
+
+	/*
+	 * Execute restore_command, which should copy
+	 * the missing WAL file from archival storage.
+	 */
+	rc = system(xlogRestoreCmd);
+
+	if (rc == 0)
+	{
+		/*
+		 * Command apparently succeeded, but let's make sure the file is
+		 * really there now and has the correct size.
+		 */
+		if (stat(xlogpath, &stat_buf) == 0)
+		{
+			if (expectedSize > 0 && stat_buf.st_size != expectedSize)
+			{
+				pg_log_error("archive file \"%s\" has wrong size: %lu instead of %lu, %s",
+							 xlogfname, (unsigned long) stat_buf.st_size,
+							 (unsigned long) expectedSize, strerror(errno));
+			}
+			else
+			{
+				xlogfd = open(xlogpath, O_RDONLY | PG_BINARY, 0);
+
+				if (xlogfd < 0)
+					pg_log_error("could not open file \"%s\" restored from archive: %s\n",
+								 xlogpath, strerror(errno));
+				else
+					return xlogfd;
+			}
+		}
+		else
+		{
+			/* Stat failed */
+			pg_log_error("could not stat file \"%s\" restored from archive: %s",
+						 xlogpath, strerror(errno));
+		}
+	}
+
+	/*
+	 * If the failure was due to any sort of signal, then it will be
+	 * misleading to return message 'could not restore file...' and
+	 * propagate result to the upper levels. We should exit right now.
+	 */
+	if (wait_result_is_any_signal(rc, false))
+		pg_fatal("restore_command failed due to the signal: %s",
+				 wait_result_to_str(rc));
+
+	pg_log_error("could not restore file \"%s\" from archive\n",
+				 xlogfname);
+	return -1;
+}
diff --git a/src/bin/pg_rewind/pg_rewind.c b/src/bin/pg_rewind/pg_rewind.c
index e75c7eed1c..72e903912e 100644
--- a/src/bin/pg_rewind/pg_rewind.c
+++ b/src/bin/pg_rewind/pg_rewind.c
@@ -53,11 +53,13 @@ int			WalSegSz;
 char	   *datadir_target = NULL;
 char	   *datadir_source = NULL;
 char	   *connstr_source = NULL;
+char	   *restore_command = NULL;
 
 static bool debug = false;
 bool		showprogress = false;
 bool		dry_run = false;
 bool		do_sync = true;
+bool		restore_wals = false;
 
 /* Target history */
 TimeLineHistoryEntry *targetHistory;
@@ -74,19 +76,22 @@ usage(const char *progname)
 	printf(_("%s resynchronizes a PostgreSQL cluster with another copy of the cluster.\n\n"), progname);
 	printf(_("Usage:\n  %s [OPTION]...\n\n"), progname);
 	printf(_("Options:\n"));
-	printf(_("  -D, --target-pgdata=DIRECTORY  existing data directory to modify\n"));
-	printf(_("      --source-pgdata=DIRECTORY  source data directory to synchronize with\n"));
-	printf(_("      --source-server=CONNSTR    source server to synchronize with\n"));
-	printf(_("  -R, --write-recovery-conf      write configuration for replication\n"
-			 "                                 (requires --source-server)\n"));
-	printf(_("  -n, --dry-run                  stop before modifying anything\n"));
-	printf(_("  -N, --no-sync                  do not wait for changes to be written\n"
-			 "                                 safely to disk\n"));
-	printf(_("  -P, --progress                 write progress messages\n"));
-	printf(_("      --no-ensure-shutdown       do not automatically fix unclean shutdown\n"));
-	printf(_("      --debug                    write a lot of debug messages\n"));
-	printf(_("  -V, --version                  output version information, then exit\n"));
-	printf(_("  -?, --help                     show this help, then exit\n"));
+	printf(_("  -D, --target-pgdata=DIRECTORY         existing data directory to modify\n"));
+	printf(_("      --source-pgdata=DIRECTORY         source data directory to synchronize with\n"));
+	printf(_("      --source-server=CONNSTR           source server to synchronize with\n"));
+	printf(_("  -R, --write-recovery-conf             write configuration for replication\n"
+			 "                                        (requires --source-server)\n"));
+	printf(_("  -c, --restore-target-wal              use restore_command in the postgresql.conf\n"));
+	printf(_("                                        to retrieve WALs from archive\n"));
+	printf(_("  -C, --target-restore-command=COMMAND  target WAL restore_command\n"));
+	printf(_("  -n, --dry-run                         stop before modifying anything\n"));
+	printf(_("  -N, --no-sync                         do not wait for changes to be written\n"
+			 "                                        safely to disk\n"));
+	printf(_("  -P, --progress                        write progress messages\n"));
+	printf(_("      --no-ensure-shutdown              do not automatically fix unclean shutdown\n"));
+	printf(_("      --debug                           write a lot of debug messages\n"));
+	printf(_("  -V, --version                         output version information, then exit\n"));
+	printf(_("  -?, --help                            show this help, then exit\n"));
 	printf(_("\nReport bugs to <pgsql-bugs@lists.postgresql.org>.\n"));
 }
 
@@ -105,6 +110,8 @@ main(int argc, char **argv)
 		{"dry-run", no_argument, NULL, 'n'},
 		{"no-sync", no_argument, NULL, 'N'},
 		{"progress", no_argument, NULL, 'P'},
+		{"restore-target-wal", no_argument, NULL, 'c'},
+		{"target-restore-command", required_argument, NULL, 'C'},
 		{"debug", no_argument, NULL, 3},
 		{NULL, 0, NULL, 0}
 	};
@@ -143,7 +150,7 @@ main(int argc, char **argv)
 		}
 	}
 
-	while ((c = getopt_long(argc, argv, "D:nNPR", long_options, &option_index)) != -1)
+	while ((c = getopt_long(argc, argv, "D:nNPRC:c", long_options, &option_index)) != -1)
 	{
 		switch (c)
 		{
@@ -155,6 +162,10 @@ main(int argc, char **argv)
 				showprogress = true;
 				break;
 
+			case 'c':
+				restore_wals = true;
+				break;
+
 			case 'n':
 				dry_run = true;
 				break;
@@ -176,6 +187,10 @@ main(int argc, char **argv)
 				datadir_target = pg_strdup(optarg);
 				break;
 
+			case 'C':
+				restore_command = pg_strdup(optarg);
+				break;
+
 			case 1:				/* --source-pgdata */
 				datadir_source = pg_strdup(optarg);
 				break;
@@ -252,6 +267,74 @@ main(int argc, char **argv)
 		exit(1);
 	}
 
+	if (restore_command != NULL)
+	{
+		if (restore_wals)
+		{
+			pg_log_error("conflicting options: both -c and -C are specified");
+			fprintf(stderr, _("You must run %s with either -c/--restore-target-wal "
+							"or -C/--target-restore-command.\n"), progname);
+			exit(1);
+		}
+
+		pg_log_debug("using command line restore_command=\'%s\'.", restore_command);
+	}
+	else if (restore_wals)
+	{
+		int   rc;
+		char  postgres_exec_path[MAXPGPATH],
+			  postgres_cmd[MAXPGPATH],
+			  cmd_output[MAXPGPATH];
+		FILE *output_fp;
+
+		/* Find postgres executable. */
+		rc = find_other_exec(argv[0], "postgres",
+							 PG_BACKEND_VERSIONSTR,
+							 postgres_exec_path);
+
+		if (rc < 0)
+		{
+			char	full_path[MAXPGPATH];
+
+			if (find_my_exec(argv[0], full_path) < 0)
+				strlcpy(full_path, progname, sizeof(full_path));
+
+			if (rc == -1)
+				pg_log_error("The program \"postgres\" is needed by %s but was not found in the\n"
+							"same directory as \"%s\".\n"
+							"Check your installation.",
+							progname, full_path);
+			else
+				pg_log_error("The program \"postgres\" was found by \"%s\"\n"
+							"but was not the same version as %s.\n"
+							"Check your installation.",
+							full_path, progname);
+			exit(1);
+		}
+
+		/* Build a command to execute for restore_command GUC retrieval if set. */
+		snprintf(postgres_cmd, sizeof(postgres_cmd), "%s -D %s -C restore_command",
+				 postgres_exec_path, datadir_target);
+
+		if ((output_fp = popen(postgres_cmd, "r")) == NULL ||
+			fgets(cmd_output, sizeof(cmd_output), output_fp) == NULL)
+			pg_fatal("could not get restore_command using %s: %s",
+					postgres_cmd, strerror(errno));
+
+		pclose(output_fp);
+
+		/* Remove trailing newline */
+		if (strchr(cmd_output, '\n') != NULL)
+			*strchr(cmd_output, '\n') = '\0';
+
+		if (!strcmp(cmd_output, ""))
+			pg_fatal("restore_command is not set on the target cluster");
+
+		restore_command = pg_strdup(cmd_output);
+
+		pg_log_debug("using config variable restore_command=\'%s\'.", restore_command);
+	}
+
 	umask(pg_mode_mask);
 
 	atexit(disconnect_atexit);
@@ -349,9 +432,8 @@ main(int argc, char **argv)
 		exit(0);
 	}
 
-	findLastCheckpoint(datadir_target, divergerec,
-					   lastcommontliIndex,
-					   &chkptrec, &chkpttli, &chkptredo);
+	findLastCheckpoint(datadir_target, divergerec, lastcommontliIndex,
+					   &chkptrec, &chkpttli, &chkptredo, restore_command);
 	pg_log_info("rewinding from last common checkpoint at %X/%X on timeline %u",
 				(uint32) (chkptrec >> 32), (uint32) chkptrec,
 				chkpttli);
@@ -377,7 +459,7 @@ main(int argc, char **argv)
 	if (showprogress)
 		pg_log_info("reading WAL in target");
 	extractPageMap(datadir_target, chkptrec, lastcommontliIndex,
-				   ControlFile_target.checkPoint);
+				   ControlFile_target.checkPoint, restore_command);
 	filemap_finalize();
 
 	if (showprogress)
diff --git a/src/bin/pg_rewind/pg_rewind.h b/src/bin/pg_rewind/pg_rewind.h
index 360b2b3da8..3abe83b38c 100644
--- a/src/bin/pg_rewind/pg_rewind.h
+++ b/src/bin/pg_rewind/pg_rewind.h
@@ -42,11 +42,13 @@ extern uint64 fetch_done;
 
 /* in parsexlog.c */
 extern void extractPageMap(const char *datadir, XLogRecPtr startpoint,
-						   int tliIndex, XLogRecPtr endpoint);
+						   int tliIndex, XLogRecPtr endpoint,
+						   const char *restoreCommand);
 extern void findLastCheckpoint(const char *datadir, XLogRecPtr searchptr,
 							   int tliIndex,
 							   XLogRecPtr *lastchkptrec, TimeLineID *lastchkpttli,
-							   XLogRecPtr *lastchkptredo);
+							   XLogRecPtr *lastchkptredo,
+							   const char *restoreCommand);
 extern XLogRecPtr readOneRecord(const char *datadir, XLogRecPtr ptr,
 								int tliIndex);
 
diff --git a/src/bin/pg_rewind/t/001_basic.pl b/src/bin/pg_rewind/t/001_basic.pl
index 95d8ccfced..73967d818d 100644
--- a/src/bin/pg_rewind/t/001_basic.pl
+++ b/src/bin/pg_rewind/t/001_basic.pl
@@ -1,7 +1,7 @@
 use strict;
 use warnings;
 use TestLib;
-use Test::More tests => 15;
+use Test::More tests => 25;
 
 use FindBin;
 use lib $FindBin::RealBin;
@@ -171,5 +171,7 @@ in master, before promotion
 # Run the test in both modes
 run_test('local');
 run_test('remote');
+run_test('archive');
+run_test('archive_conf');
 
 exit(0);
diff --git a/src/bin/pg_rewind/t/002_databases.pl b/src/bin/pg_rewind/t/002_databases.pl
index 1db534c0dc..d046a755e5 100644
--- a/src/bin/pg_rewind/t/002_databases.pl
+++ b/src/bin/pg_rewind/t/002_databases.pl
@@ -1,7 +1,7 @@
 use strict;
 use warnings;
 use TestLib;
-use Test::More tests => 7;
+use Test::More tests => 13;
 
 use FindBin;
 use lib $FindBin::RealBin;
@@ -70,5 +70,7 @@ template1
 # Run the test in both modes.
 run_test('local');
 run_test('remote');
+run_test('archive');
+run_test('archive_conf');
 
 exit(0);
diff --git a/src/bin/pg_rewind/t/003_extrafiles.pl b/src/bin/pg_rewind/t/003_extrafiles.pl
index f4710440fc..ab1fbbc125 100644
--- a/src/bin/pg_rewind/t/003_extrafiles.pl
+++ b/src/bin/pg_rewind/t/003_extrafiles.pl
@@ -3,7 +3,7 @@
 use strict;
 use warnings;
 use TestLib;
-use Test::More tests => 5;
+use Test::More tests => 9;
 
 use File::Find;
 
@@ -90,5 +90,7 @@ sub run_test
 # Run the test in both modes.
 run_test('local');
 run_test('remote');
+run_test('archive');
+run_test('archive_conf');
 
 exit(0);
diff --git a/src/bin/pg_rewind/t/RewindTest.pm b/src/bin/pg_rewind/t/RewindTest.pm
index 82fa220ac8..58055bb494 100644
--- a/src/bin/pg_rewind/t/RewindTest.pm
+++ b/src/bin/pg_rewind/t/RewindTest.pm
@@ -35,7 +35,9 @@ use Carp;
 use Config;
 use Exporter 'import';
 use File::Copy;
-use File::Path qw(rmtree);
+use File::Glob ':bsd_glob';
+use File::Path qw(remove_tree make_path);
+use File::Spec::Functions qw(catdir catfile);
 use IPC::Run qw(run);
 use PostgresNode;
 use TestLib;
@@ -216,6 +218,38 @@ sub promote_standby
 	return;
 }
 
+# Moves WAL files to the temporary location and returns restore_command
+# to get them back.
+sub move_wals
+{
+	my $tmp_dir = shift;
+	my $master_pgdata = shift;
+	my $wals_archive_dir = catdir($tmp_dir, "master_wals_archive");
+	my @wal_files = bsd_glob catfile($master_pgdata, "pg_wal", "0000000*");
+	my $restore_command;
+
+	remove_tree($wals_archive_dir);
+	make_path($wals_archive_dir) or die;
+
+	# Move all old master WAL files to the archive.
+	# Old master should be stopped at this point.
+	foreach my $wal_file (@wal_files)
+	{
+		move($wal_file, "$wals_archive_dir") or die;
+	}
+
+	if ($windows_os)
+	{
+		$restore_command = "copy $wals_archive_dir\\\%f \%p";
+	}
+	else
+	{
+		$restore_command = "cp $wals_archive_dir/\%f \%p";
+	}
+
+	return $restore_command;
+}
+
 sub run_pg_rewind
 {
 	my $test_mode       = shift;
@@ -227,10 +261,23 @@ sub run_pg_rewind
 	# Append the rewind-specific role to the connection string.
 	$standby_connstr = "$standby_connstr user=rewind_user";
 
-	# Stop the master and be ready to perform the rewind.  The cluster
-	# needs recovery to finish once, and pg_rewind makes sure that it
-	# happens automatically.
-	$node_master->stop('immediate');
+	if ($test_mode eq 'archive' || $test_mode eq 'archive_conf')
+	{
+		# We test pg_rewind with restore_command by simply moving all
+		# WAL to another location.  It leads to failed ensureCleanShutdown
+		# execution.  Since it is difficult to emulate a situation, when
+		# keeping the last WAL segment is enough for startup recovery, but
+		# not enough for successful pg_rewind run, we run these modes with
+		# --no-ensure-shutdown.  So stop the master gracefully.
+		$node_master->stop;
+	}
+	else
+	{
+		# Stop the master and be ready to perform the rewind.  The cluster
+		# needs recovery to finish once, and pg_rewind makes sure that it
+		# happens automatically.
+		$node_master->stop('immediate');
+	}
 
 	# At this point, the rewind processing is ready to run.
 	# We now have a very simple scenario with a few diverged WAL record.
@@ -284,6 +331,54 @@ sub run_pg_rewind
 		$node_standby->safe_psql('postgres',
 			"ALTER ROLE rewind_user WITH REPLICATION;");
 	}
+	elsif ($test_mode eq "archive")
+	{
+
+		# Do rewind using a local pgdata as source and
+		# specified directory with target WALs archive.
+		my $restore_command = move_wals($tmp_folder, $master_pgdata);
+
+		# Stop the new master and be ready to perform the rewind.
+		$node_standby->stop;
+
+		command_ok(
+			[
+				'pg_rewind',
+				"--debug",
+				"--source-pgdata=$standby_pgdata",
+				"--target-pgdata=$master_pgdata",
+				"--no-sync", "--no-ensure-shutdown",
+				"--target-restore-command=$restore_command"
+			],
+			'pg_rewind archive');
+	}
+	elsif ($test_mode eq "archive_conf")
+	{
+
+		# Do rewind using a local pgdata as source and
+		# specified directory with target WALs archive.
+		my $master_conf_path = catfile($master_pgdata, 'postgresql.conf');
+		my $restore_command = move_wals($tmp_folder, $master_pgdata);
+
+		# Stop the new master and be ready to perform the rewind.
+		$node_standby->stop;
+
+		# Add restore_command to postgresql.conf of target cluster.
+		open(my $conf_fd, ">>", $master_conf_path) or die;
+		print $conf_fd "\nrestore_command='$restore_command'";
+		close $conf_fd;
+
+		command_ok(
+			[
+				'pg_rewind',
+				"--debug",
+				"--source-pgdata=$standby_pgdata",
+				"--target-pgdata=$master_pgdata",
+				"--no-sync", "--no-ensure-shutdown",
+				"-c"
+			],
+			'pg_rewind archive_conf');
+	}
 	else
 	{
 

base-commit: 88d45ac752ae49dbfafeb163b07381d3b8a6b601
-- 
2.17.1

