From f14f6a2419b06a6e280e28e30a005fbbc375456e Mon Sep 17 00:00:00 2001
From: Polina Bungina <bungina@gmail.com>
Date: Sun, 21 Aug 2022 17:48:32 +0200
Subject: [PATCH v1] Be more picky with WAL segment deletion in pg_rewind

Make pg_rewind to be a bit wiser in terms of creating filemap:
preserve on the target all WAL segments that contain records between the
last common checkpoint and the point of divergence.

Co-authored-by: Alexander Kukushkin <cyberdemn@gmail.com>
---
 src/bin/pg_rewind/filemap.c   | 50 +++++++++++++++++++++++++++++++++++
 src/bin/pg_rewind/filemap.h   |  3 +++
 src/bin/pg_rewind/parsexlog.c | 10 +++++++
 src/bin/pg_rewind/pg_rewind.c |  2 ++
 4 files changed, 65 insertions(+)

diff --git a/src/bin/pg_rewind/filemap.c b/src/bin/pg_rewind/filemap.c
index 269ed6446e6..9633c0dc926 100644
--- a/src/bin/pg_rewind/filemap.c
+++ b/src/bin/pg_rewind/filemap.c
@@ -64,6 +64,27 @@ static file_entry_t *lookup_filehash_entry(const char *path);
 static int	final_filemap_cmp(const void *a, const void *b);
 static bool check_file_excluded(const char *path, bool is_source);
 
+typedef struct skipwal_t {
+	const char *path;
+	uint32		status;
+} skipwal_t;
+
+#define SH_PREFIX		keepwalhash
+#define SH_ELEMENT_TYPE	skipwal_t
+#define SH_KEY_TYPE		const char *
+#define	SH_KEY			path
+#define SH_HASH_KEY(tb, key)	hash_string_pointer(key)
+#define SH_EQUAL(tb, a, b)		(strcmp(a, b) == 0)
+#define	SH_SCOPE		static inline
+#define SH_RAW_ALLOCATOR	pg_malloc0
+#define SH_DECLARE
+#define SH_DEFINE
+#include "lib/simplehash.h"
+
+static keepwalhash_hash *keepwalhash = NULL;
+
+static bool keepwalhash_entry_exists(const char *path);
+
 /*
  * Definition of one element part of an exclusion list, used to exclude
  * contents when rewinding.  "name" is the name of the file or path to
@@ -207,6 +228,33 @@ lookup_filehash_entry(const char *path)
 	return filehash_lookup(filehash, path);
 }
 
+void
+keepwalhash_init(void)
+{
+	keepwalhash = keepwalhash_create(FILEHASH_INITIAL_SIZE, NULL);
+}
+
+void
+insert_keepwalhash_entry(const char *path)
+{
+	skipwal_t   *entry;
+	bool		found;
+
+	/* Should only be called with keepwalhash initialized */
+	Assert(keepwalhash);
+
+	entry = keepwalhash_insert(keepwalhash, path, &found);
+
+	if (!found)
+		entry->path = pg_strdup(path);
+}
+
+static bool
+keepwalhash_entry_exists(const char *path)
+{
+	return keepwalhash_lookup(keepwalhash, path) != NULL;
+}
+
 /*
  * Callback for processing source file list.
  *
@@ -658,6 +706,8 @@ decide_file_action(file_entry_t *entry)
 			return FILE_ACTION_NONE;
 	}
 
+	if (keepwalhash_entry_exists(path))
+		return FILE_ACTION_NONE;
 	/*
 	 * Handle cases where the file is missing from one of the systems.
 	 */
diff --git a/src/bin/pg_rewind/filemap.h b/src/bin/pg_rewind/filemap.h
index 0e011fbb0b3..c637a5042d9 100644
--- a/src/bin/pg_rewind/filemap.h
+++ b/src/bin/pg_rewind/filemap.h
@@ -110,4 +110,7 @@ extern filemap_t *decide_file_actions(void);
 extern void calculate_totals(filemap_t *filemap);
 extern void print_filemap(filemap_t *filemap);
 
+extern void keepwalhash_init(void);
+extern void insert_keepwalhash_entry(const char *path);
+
 #endif							/* FILEMAP_H */
diff --git a/src/bin/pg_rewind/parsexlog.c b/src/bin/pg_rewind/parsexlog.c
index 53f011a2fe2..b872f782f11 100644
--- a/src/bin/pg_rewind/parsexlog.c
+++ b/src/bin/pg_rewind/parsexlog.c
@@ -48,6 +48,7 @@ typedef struct XLogPageReadPrivate
 {
 	const char *restoreCommand;
 	int			tliIndex;
+	bool		keepWalSeg;
 } XLogPageReadPrivate;
 
 static int	SimpleXLogPageRead(XLogReaderState *xlogreader,
@@ -73,6 +74,7 @@ extractPageMap(const char *datadir, XLogRecPtr startpoint, int tliIndex,
 
 	private.tliIndex = tliIndex;
 	private.restoreCommand = restoreCommand;
+	private.keepWalSeg = false;
 	xlogreader = XLogReaderAllocate(WalSegSz, datadir,
 									XL_ROUTINE(.page_read = &SimpleXLogPageRead),
 									&private);
@@ -132,6 +134,7 @@ readOneRecord(const char *datadir, XLogRecPtr ptr, int tliIndex,
 
 	private.tliIndex = tliIndex;
 	private.restoreCommand = restoreCommand;
+	private.keepWalSeg = false;
 	xlogreader = XLogReaderAllocate(WalSegSz, datadir,
 									XL_ROUTINE(.page_read = &SimpleXLogPageRead),
 									&private);
@@ -192,6 +195,7 @@ findLastCheckpoint(const char *datadir, XLogRecPtr forkptr, int tliIndex,
 
 	private.tliIndex = tliIndex;
 	private.restoreCommand = restoreCommand;
+	private.keepWalSeg = true;
 	xlogreader = XLogReaderAllocate(WalSegSz, datadir,
 									XL_ROUTINE(.page_read = &SimpleXLogPageRead),
 									&private);
@@ -297,6 +301,12 @@ SimpleXLogPageRead(XLogReaderState *xlogreader, XLogRecPtr targetPagePtr,
 		XLogFileName(xlogfname, targetHistory[private->tliIndex].tli,
 					 xlogreadsegno, WalSegSz);
 
+		if (private->keepWalSeg)
+		{
+			snprintf(xlogfpath, MAXPGPATH, XLOGDIR "/%s", xlogfname);
+			insert_keepwalhash_entry(xlogfpath);
+		}
+
 		snprintf(xlogfpath, MAXPGPATH, "%s/" XLOGDIR "/%s",
 				 xlogreader->segcxt.ws_dir, xlogfname);
 
diff --git a/src/bin/pg_rewind/pg_rewind.c b/src/bin/pg_rewind/pg_rewind.c
index 1ff8da16764..85e087a7ac0 100644
--- a/src/bin/pg_rewind/pg_rewind.c
+++ b/src/bin/pg_rewind/pg_rewind.c
@@ -405,6 +405,8 @@ main(int argc, char **argv)
 		exit(0);
 	}
 
+	keepwalhash_init();
+
 	findLastCheckpoint(datadir_target, divergerec, lastcommontliIndex,
 					   &chkptrec, &chkpttli, &chkptredo, restore_command);
 	pg_log_info("rewinding from last common checkpoint at %X/%X on timeline %u",
-- 
2.36.1

