From 279affc0b822f41582b9f5ddfcaf07e554f87e08 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddyforpostgres@gmail.com>
Date: Fri, 21 Oct 2022 04:53:21 +0000
Subject: [PATCH v9] Avoid memory leaks during backups using SQL-callable
 functions

Any failure while taking backups using SQL-callable functions can
leak the memory allocated for backup_state and tablespace_map, as
they both need to be long-lived, they are being created in
TopMemoryContext.

To fix the memory leak problem, we create a special session-level
memory context as a direct child of TopMemoryContext so that the
memory allocated is carried across. We delete the memory context at
the end of pg_backup_stop(). We keep the memory allocated in this
memory context less, because any error before reaching
pg_backup_stop() can still leak the memory until pg_backup_start()
is called again. While this is not smart, it helps to keep things
simple.

Author: Bharath Rupireddy
Reviewed-by: Robert Haas, Alvaro Herrera
Reviewed-by: Cary Huang, Michael Paquier
Discussion: https://www.postgresql.org/message-id/CALj2ACXqvfKF2B0beQ=aJMdWnpNohmBPsRg=EDQj_6y1t2O8mQ@mail.gmail.com
---
 src/backend/access/transam/xlogfuncs.c | 41 ++++++++++++++------------
 1 file changed, 22 insertions(+), 19 deletions(-)

diff --git a/src/backend/access/transam/xlogfuncs.c b/src/backend/access/transam/xlogfuncs.c
index a801a94fe8..6663c4cc8f 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 long-lived workspace for SQL-callable backup functions. */
+static MemoryContext backupcontext = NULL;
+
 /*
  * pg_backup_start: set up for taking an on-line backup dump
  *
@@ -72,27 +75,26 @@ pg_backup_start(PG_FUNCTION_ARGS)
 
 	/*
 	 * backup_state and tablespace_map need to be long-lived as they are used
-	 * in pg_backup_stop().
+	 * in pg_backup_stop(). Create a special session-level memory context as a
+	 * direct child of TopMemoryContext so that the memory allocated is carried
+	 * across. We keep the memory allocated in this memory context less,
+	 * because any error before reaching pg_backup_stop() can leak the memory
+	 * until pg_backup_start() is called again. While this is not smart, it
+	 * helps to keep things simple.
 	 */
-	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
-
-	/* Allocate backup state or reset it, if it comes from a previous run */
-	if (backup_state == NULL)
-		backup_state = (BackupState *) palloc0(sizeof(BackupState));
+	if (backupcontext == NULL)
+		backupcontext = AllocSetContextCreate(TopMemoryContext,
+											  "on-line backup context",
+											  ALLOCSET_START_SMALL_SIZES);
 	else
-		MemSet(backup_state, 0, sizeof(BackupState));
-
-	/*
-	 * tablespace_map may have been created in a previous backup, so take this
-	 * occasion to clean it.
-	 */
-	if (tablespace_map != NULL)
 	{
-		pfree(tablespace_map->data);
-		pfree(tablespace_map);
+		MemoryContextReset(backupcontext);
+		backup_state = NULL;
 		tablespace_map = NULL;
 	}
 
+	oldcontext = MemoryContextSwitchTo(backupcontext);
+	backup_state = (BackupState *) palloc0(sizeof(BackupState));
 	tablespace_map = makeStringInfo();
 	MemoryContextSwitchTo(oldcontext);
 
@@ -157,12 +159,13 @@ pg_backup_stop(PG_FUNCTION_ARGS)
 	values[2] = CStringGetTextDatum(tablespace_map->data);
 
 	/* Deallocate backup-related variables */
-	pfree(backup_state);
+	pfree(backup_label);
+
+	/* Clean up the session-level backup memory context */
+	MemoryContextDelete(backupcontext);
 	backup_state = NULL;
-	pfree(tablespace_map->data);
-	pfree(tablespace_map);
 	tablespace_map = NULL;
-	pfree(backup_label);
+	backupcontext = NULL;
 
 	/* Returns the record as Datum */
 	PG_RETURN_DATUM(HeapTupleGetDatum(heap_form_tuple(tupdesc, values, nulls)));
-- 
2.34.1

