From 36607bb46a1328a8b4d63169eb1739cf3e6769fd Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddyforpostgres@gmail.com>
Date: Wed, 28 Sep 2022 09:04:00 +0000
Subject: [PATCH v2] Avoid memory leaks during backups

Postgres currently can leak memory if a failure occurs during base
backup in do_pg_backup_start() or do_pg_backup_stop() or
perform_base_backup(). The palloc'd memory such as backup_state or
tablespace_map in xlogfuncs.c or basebackup.c or tablespaceinfo in
perform_base_backup() or any other, is left-out which may cause
memory bloat on the server eventually. To experience this issue,
run pg_basebackup with --label name longer than 1024 characters and
observe memory with watch command, the memory usage goes up.

It looks like the memory leak issue has been there for quite some
time.

To solve the above problem, leverage PG_TRY()-PG_FINALLY()-PG_END_TRY()
mechanism and memory context. The design of the patch is as follows:

For backup functions or on-line backups, a new memory context is
created as a child of TopMemoryContext in pg_backup_start() which
gets deleted upon ERROR or at the end of pg_backup_stop().

For perform_base_backup() or base backups, a new memory context is
created as a child of CurrentMemoryContext which gets deleted upon
ERROR or at the end of perform_base_backup().
---
 src/backend/access/transam/xlogfuncs.c | 114 ++++++++++++++++++-------
 src/backend/backup/basebackup.c        |  41 ++++++++-
 src/backend/utils/error/elog.c         |  17 ++++
 src/include/utils/elog.h               |   1 +
 4 files changed, 141 insertions(+), 32 deletions(-)

diff --git a/src/backend/access/transam/xlogfuncs.c b/src/backend/access/transam/xlogfuncs.c
index a801a94fe8..6c13da91bd 100644
--- a/src/backend/access/transam/xlogfuncs.c
+++ b/src/backend/access/transam/xlogfuncs.c
@@ -45,6 +45,9 @@
 static BackupState *backup_state = NULL;
 static StringInfo tablespace_map = NULL;
 
+/* A workspace for on-line backup. */
+static MemoryContext backupcontext = NULL;
+
 /*
  * pg_backup_start: set up for taking an on-line backup dump
  *
@@ -71,33 +74,52 @@ pg_backup_start(PG_FUNCTION_ARGS)
 				 errmsg("a backup is already in progress in this session")));
 
 	/*
-	 * backup_state and tablespace_map need to be long-lived as they are used
-	 * in pg_backup_stop().
+	 * Start the on-line backup, but make sure we clean up the allocated memory
+	 * even if an error occurs.
 	 */
-	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
+	PG_TRY();
+	{
+		/*
+		 * backup_state and tablespace_map need to be long-lived as they are
+		 * used in pg_backup_stop(). Create a special memory context as a
+		 * direct child of TopMemoryContext so that the memory allocated is
+		 * cleaned in case of errors.
+		 */
+		if (backupcontext == NULL)
+		{
+			backupcontext = AllocSetContextCreate(TopMemoryContext,
+												  "on-line backup context",
+												  ALLOCSET_DEFAULT_SIZES);
+		}
 
-	/* Allocate backup state or reset it, if it comes from a previous run */
-	if (backup_state == NULL)
+		oldcontext = MemoryContextSwitchTo(backupcontext);
+
+		Assert(backup_state == NULL);
+		Assert(tablespace_map == NULL);
 		backup_state = (BackupState *) palloc0(sizeof(BackupState));
-	else
-		MemSet(backup_state, 0, sizeof(BackupState));
+		tablespace_map = makeStringInfo();
 
-	/*
-	 * tablespace_map may have been created in a previous backup, so take this
-	 * occasion to clean it.
-	 */
-	if (tablespace_map != NULL)
+		register_persistent_abort_backup_handler();
+		do_pg_backup_start(backupidstr, fast, NULL, backup_state, tablespace_map);
+
+		MemoryContextSwitchTo(oldcontext);
+	}
+	PG_FINALLY();
 	{
-		pfree(tablespace_map->data);
-		pfree(tablespace_map);
-		tablespace_map = NULL;
+		/*
+		 * Delete on-line backup memory context only when an ERROR occurred,
+		 * because the memory allocated in backupcontext is used in
+		 * pg_backup_stop() for non-ERROR cases.
+		 */
+		if (backupcontext != NULL && geterrlevel() == ERROR)
+		{
+			MemoryContextDelete(backupcontext);
+			backupcontext = NULL;
+			backup_state = NULL;
+			tablespace_map = NULL;
+		}
 	}
-
-	tablespace_map = makeStringInfo();
-	MemoryContextSwitchTo(oldcontext);
-
-	register_persistent_abort_backup_handler();
-	do_pg_backup_start(backupidstr, fast, NULL, backup_state, tablespace_map);
+	PG_END_TRY();
 
 	PG_RETURN_LSN(backup_state->startpoint);
 }
@@ -132,6 +154,7 @@ pg_backup_stop(PG_FUNCTION_ARGS)
 	bool		waitforarchive = PG_GETARG_BOOL(0);
 	char	   *backup_label;
 	SessionBackupState status = get_backup_status();
+	MemoryContext oldcontext;
 
 	/* Initialize attributes information in the tuple descriptor */
 	if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
@@ -145,24 +168,55 @@ pg_backup_stop(PG_FUNCTION_ARGS)
 
 	Assert(backup_state != NULL);
 	Assert(tablespace_map != NULL);
+	Assert(backupcontext != NULL);
 
-	/* Stop the backup */
-	do_pg_backup_stop(backup_state, waitforarchive);
+	/*
+	 * Stop the on-line backup, but make sure we clean up the allocated memory
+	 * even if an error occurs.
+	 */
+	PG_TRY();
+	{
+		oldcontext = MemoryContextSwitchTo(backupcontext);
 
-	/* Build the contents of backup_label */
-	backup_label = build_backup_content(backup_state, false);
+		do_pg_backup_stop(backup_state, waitforarchive);
+
+		/* Build the contents of backup_label */
+		backup_label = build_backup_content(backup_state, false);
+
+		/*
+		 * Switch back from on-line backup memory context so that the results
+		 * are in right memory context.
+		 */
+		MemoryContextSwitchTo(oldcontext);
+	}
+	PG_FINALLY();
+	{
+		/*
+		 * Delete on-line backup memory context only when an ERROR occurred,
+		 * because the memory allocated in backupcontext is used in further
+		 * down in this function for non-ERROR cases.
+		 */
+		if (backupcontext != NULL && geterrlevel() == ERROR)
+		{
+			MemoryContextDelete(backupcontext);
+			backupcontext = NULL;
+			backup_label = NULL;
+			backup_state = NULL;
+			tablespace_map = NULL;
+		}
+	}
+	PG_END_TRY();
 
 	values[0] = LSNGetDatum(backup_state->stoppoint);
 	values[1] = CStringGetTextDatum(backup_label);
 	values[2] = CStringGetTextDatum(tablespace_map->data);
 
-	/* Deallocate backup-related variables */
-	pfree(backup_state);
+	/* Delete on-line backup memory context */
+	MemoryContextDelete(backupcontext);
+	backupcontext = NULL;
+	backup_label = NULL;
 	backup_state = NULL;
-	pfree(tablespace_map->data);
-	pfree(tablespace_map);
 	tablespace_map = NULL;
-	pfree(backup_label);
 
 	/* Returns the record as Datum */
 	PG_RETURN_DATUM(HeapTupleGetDatum(heap_form_tuple(tupdesc, values, nulls)));
diff --git a/src/backend/backup/basebackup.c b/src/backend/backup/basebackup.c
index 1434bcdd85..60f216b891 100644
--- a/src/backend/backup/basebackup.c
+++ b/src/backend/backup/basebackup.c
@@ -42,6 +42,7 @@
 #include "storage/reinit.h"
 #include "utils/builtins.h"
 #include "utils/guc.h"
+#include "utils/memutils.h"
 #include "utils/ps_status.h"
 #include "utils/relcache.h"
 #include "utils/resowner.h"
@@ -75,6 +76,9 @@ typedef struct
 	pg_checksum_type manifest_checksum_type;
 } basebackup_options;
 
+/* A workspace for base backup. */
+static MemoryContext backupcontext = NULL;
+
 static int64 sendTablespace(bbsink *sink, char *path, char *spcoid, bool sizeonly,
 							struct backup_manifest_info *manifest);
 static int64 sendDir(bbsink *sink, const char *path, int basepathlen, bool sizeonly,
@@ -235,6 +239,7 @@ perform_base_backup(basebackup_options *opt, bbsink *sink)
 	backup_manifest_info manifest;
 	BackupState *backup_state;
 	StringInfo	tablespace_map;
+	MemoryContext oldcontext;
 
 	/* Initial backup state, insofar as we know it now. */
 	state.tablespaces = NIL;
@@ -252,6 +257,14 @@ perform_base_backup(basebackup_options *opt, bbsink *sink)
 	InitializeBackupManifest(&manifest, opt->manifest,
 							 opt->manifest_checksum_type);
 
+	/*
+	 * Switch to base backup memory context. We do this only after
+	 * InitializeBackupManifest() as it uses BufFile which needs to be in the
+	 * ResourceOwner's memory context.
+	 */
+	Assert(backupcontext != NULL);
+	oldcontext = MemoryContextSwitchTo(backupcontext);
+
 	total_checksum_failures = 0;
 
 	/* Allocate backup related varilables. */
@@ -660,6 +673,8 @@ perform_base_backup(basebackup_options *opt, bbsink *sink)
 	 */
 	FreeBackupManifest(&manifest);
 
+	MemoryContextSwitchTo(oldcontext);
+
 	/* clean up the resource owner we created */
 	WalSndResourceCleanup(true);
 
@@ -1011,16 +1026,38 @@ SendBaseBackup(BaseBackupCmd *cmd)
 	sink = bbsink_progress_new(sink, opt.progress);
 
 	/*
-	 * Perform the base backup, but make sure we clean up the bbsink even if
-	 * an error occurs.
+	 * Perform the base backup, but make sure we clean up the bbsink and
+	 * allocated memory even if an error occurs.
 	 */
 	PG_TRY();
 	{
+		/*
+		 * Perform the base backup in a special memory context to not leak any
+		 * memory in case of errors. perform_base_backup() will switch to
+		 * backupcontext in the appropriate time.
+		 */
+		if (backupcontext == NULL)
+		{
+			backupcontext = AllocSetContextCreate(CurrentMemoryContext,
+												  "base backup context",
+												  ALLOCSET_DEFAULT_SIZES);
+		}
+
 		perform_base_backup(&opt, sink);
 	}
 	PG_FINALLY();
 	{
 		bbsink_cleanup(sink);
+
+		/*
+		 * Delete base backup memory context. perform_base_backup() would have
+		 * switched from backupcontext to the right memory context at the end.
+		 */
+		if (backupcontext != NULL)
+		{
+			MemoryContextDelete(backupcontext);
+			backupcontext = NULL;
+		}
 	}
 	PG_END_TRY();
 }
diff --git a/src/backend/utils/error/elog.c b/src/backend/utils/error/elog.c
index eb724a9d7f..18710d2d89 100644
--- a/src/backend/utils/error/elog.c
+++ b/src/backend/utils/error/elog.c
@@ -1404,6 +1404,23 @@ geterrcode(void)
 	return edata->sqlerrcode;
 }
 
+/*
+ * geterrlevel --- return the currently error level
+ */
+int
+geterrlevel(void)
+{
+	ErrorData  *edata;
+
+	/* Exit if errstart was not called at all */
+	if (errordata_stack_depth < 0)
+		return 0;
+
+	edata = &errordata[errordata_stack_depth];
+
+	return edata->elevel;
+}
+
 /*
  * geterrposition --- return the currently set error position (0 if none)
  *
diff --git a/src/include/utils/elog.h b/src/include/utils/elog.h
index 4dd9658a3c..a8955f9572 100644
--- a/src/include/utils/elog.h
+++ b/src/include/utils/elog.h
@@ -222,6 +222,7 @@ extern int	internalerrquery(const char *query);
 extern int	err_generic_string(int field, const char *str);
 
 extern int	geterrcode(void);
+extern int	geterrlevel(void);
 extern int	geterrposition(void);
 extern int	getinternalerrposition(void);
 
-- 
2.34.1

