Fix Subtransaction memory leak

src/backend/distributed/commands/variableset.c

@@ -127,7 +127,8 @@ PostprocessVariableSetStmt(VariableSetStmt *setStmt, const char *setStmtString)
 	/* haven't seen any SET stmts so far in this (sub-)xact: initialize StringInfo */
 	if (activeSetStmts == NULL)
 	{
-		MemoryContext old_context = MemoryContextSwitchTo(CurTransactionContext);
+		/* see comments in PushSubXact on why we allocate this in TopTransactionContext */
+		MemoryContext old_context = MemoryContextSwitchTo(TopTransactionContext);
 		activeSetStmts = makeStringInfo();
 		MemoryContextSwitchTo(old_context);
 	}

src/backend/distributed/test/xact_stats.c

@@ -0,0 +1,56 @@
+/*-------------------------------------------------------------------------
+ *
+ * xact_stats.c
+ *
+ * This file contains functions to provide helper UDFs for testing transaction
+ * statistics.
+ *
+ * Copyright (c) Citus Data, Inc.
+ *
+ *-------------------------------------------------------------------------
+ */
+#include <sys/stat.h>
+#include <unistd.h>
+
+#include "postgres.h"
+#include "funcapi.h"
+#include "libpq-fe.h"
+#include "miscadmin.h"
+#include "pgstat.h"
+
+static Size MemoryContextTotalSpace(MemoryContext context);
+
+PG_FUNCTION_INFO_V1(top_transaction_context_size);
+
+/*
+ * top_transaction_context_size returns current size of TopTransactionContext.
+ */
+Datum
+top_transaction_context_size(PG_FUNCTION_ARGS)
+{
+	Size totalSpace = MemoryContextTotalSpace(TopTransactionContext);
+	PG_RETURN_INT64(totalSpace);
+}
+
+
+/*
+ * MemoryContextTotalSpace returns total space allocated in context and its children.
+ */
+static Size
+MemoryContextTotalSpace(MemoryContext context)
+{
+	Size totalSpace = 0;
+
+	MemoryContextCounters totals = { 0 };
+	TopTransactionContext->methods->stats(TopTransactionContext, NULL, NULL, &totals);
+	totalSpace += totals.totalspace;
+
+	for (MemoryContext child = context->firstchild;
+		 child != NULL;
+		 child = child->nextchild)
+	{
+		totalSpace += MemoryContextTotalSpace(child);
+	}
+
+	return totalSpace;
+}

src/backend/distributed/transaction/transaction_management.c

@@ -593,7 +593,17 @@ AdjustMaxPreparedTransactions(void)
 static void
 PushSubXact(SubTransactionId subId)
 {
-	MemoryContext old_context = MemoryContextSwitchTo(CurTransactionContext);
+	/*
+	 * We need to allocate these in TopTransactionContext instead of current
+	 * subxact's memory context. This is because AtSubCommit_Memory won't
+	 * delete the subxact's memory context unless it is empty, and this
+	 * can cause in memory leaks. For emptiness it just checks if the memory
+	 * has been reset, and we cannot reset the subxact context since other
+	 * data can be in the context that are needed by upper commits.
+	 *
+	 * See https://github.com/citusdata/citus/issues/3999
+	 */
+	MemoryContext old_context = MemoryContextSwitchTo(TopTransactionContext);
 
 	/* save provided subId as well as propagated SET LOCAL stmts */
 	SubXactContext *state = palloc(sizeof(SubXactContext));
@@ -612,19 +622,34 @@ PushSubXact(SubTransactionId subId)
 static void
 PopSubXact(SubTransactionId subId)
 {
-	MemoryContext old_context = MemoryContextSwitchTo(CurTransactionContext);
 	SubXactContext *state = linitial(activeSubXactContexts);
 
-	/*
-	 * the previous activeSetStmts is already invalid because it's in the now-
-	 * aborted subxact (what we're popping), so no need to free before assign-
-	 * ing with the setLocalCmds of the popped context
-	 */
 	Assert(state->subId == subId);
-	activeSetStmts = state->setLocalCmds;
-	activeSubXactContexts = list_delete_first(activeSubXactContexts);
 
-	MemoryContextSwitchTo(old_context);
+	/*
+	 * Free activeSetStmts to avoid memory leaks when we create subxacts
+	 * for each row, e.g. in exception handling of UDFs.
+	 */
+	if (activeSetStmts != NULL)
+	{
+		pfree(activeSetStmts->data);
+		pfree(activeSetStmts);
+	}
+
+	/*
+	 * SET LOCAL commands are local to subxact blocks. When a subxact commits
+	 * or rolls back, we should roll back our set of SET LOCAL commands to the
+	 * ones we had in the upper commit.
+	 */
+	activeSetStmts = state->setLocalCmds;
+
+	/*
+	 * Free state to avoid memory leaks when we create subxacts for each row,
+	 * e.g. in exception handling of UDFs.
+	 */
+	pfree(state);
+
+	activeSubXactContexts = list_delete_first(activeSubXactContexts);
 }
 
 

src/test/regress/expected/multi_subtransactions.out

@@ -1,3 +1,5 @@
+CREATE SCHEMA multi_subtransactions;
+SET search_path TO 'multi_subtransactions';
 CREATE TABLE artists (
     id bigint NOT NULL,
     name text NOT NULL
@@ -159,7 +161,7 @@ BEGIN;
 DELETE FROM artists;
 SAVEPOINT s1;
 INSERT INTO artists SELECT NULL, NULL FROM generate_series(1, 5) i;
-ERROR:  the partition column of table public.artists cannot be NULL
+ERROR:  the partition column of table multi_subtransactions.artists cannot be NULL
 ROLLBACK TO s1;
 INSERT INTO artists VALUES (11, 'Egon Schiele');
 COMMIT;
@@ -175,7 +177,7 @@ BEGIN;
 DELETE FROM artists;
 SAVEPOINT s1;
 INSERT INTO artists(name) SELECT 'a' FROM generate_series(1, 5) i;
-ERROR:  the partition column of table public.artists should have a value
+ERROR:  the partition column of table multi_subtransactions.artists should have a value
 ROLLBACK TO s1;
 INSERT INTO artists VALUES (12, 'Marc Chagall');
 COMMIT;
@@ -401,6 +403,29 @@ SELECT * FROM researchers WHERE lab_id=10;
  32 |     10 | Raymond Smullyan
 (2 rows)
 
+-- Verify that we don't have a memory leak in subtransactions
+-- See https://github.com/citusdata/citus/pull/4000
+CREATE FUNCTION text2number(v_value text) RETURNS numeric
+    LANGUAGE plpgsql VOLATILE
+    AS $$
+BEGIN
+ RETURN v_value::numeric;
+exception
+    when others then
+        return null;
+END;
+$$;
+-- if we leak at least an integer in each subxact, then size of TopTransactionSize
+-- will be way beyond the 50k limit. If issue #3999 happens, then this will also take
+-- a long time, since for each row we will create a memory context that is not destroyed
+-- until the end of command.
+SELECT max(text2number('1234')), max(public.top_transaction_context_size()) > 50000 AS leaked
+FROM generate_series(1, 20000);
+ max  | leaked
+---------------------------------------------------------------------
+ 1234 | f
+(1 row)
+
 -- Clean-up
-DROP TABLE artists;
-DROP TABLE researchers;
+SET client_min_messages TO ERROR;
+DROP SCHEMA multi_subtransactions CASCADE;

src/test/regress/expected/multi_test_helpers_superuser.out

@@ -49,3 +49,7 @@ CREATE OR REPLACE FUNCTION pg_catalog.partition_task_list_results(resultIdPrefix
                   targetShardIndex int)
     LANGUAGE C STRICT VOLATILE
     AS 'citus', $$partition_task_list_results$$;
+-- get size of TopTransactionContext
+CREATE OR REPLACE FUNCTION top_transaction_context_size() RETURNS BIGINT
+LANGUAGE C STRICT VOLATILE
+AS 'citus', $$top_transaction_context_size$$;

src/test/regress/sql/multi_subtransactions.sql

@@ -1,3 +1,6 @@
+CREATE SCHEMA multi_subtransactions;
+SET search_path TO 'multi_subtransactions';
+
 CREATE TABLE artists (
     id bigint NOT NULL,
     name text NOT NULL
@@ -305,7 +308,27 @@ COMMIT;
 
 SELECT * FROM researchers WHERE lab_id=10;
 
--- Clean-up
-DROP TABLE artists;
-DROP TABLE researchers;
+-- Verify that we don't have a memory leak in subtransactions
+-- See https://github.com/citusdata/citus/pull/4000
 
+CREATE FUNCTION text2number(v_value text) RETURNS numeric
+    LANGUAGE plpgsql VOLATILE
+    AS $$
+BEGIN
+ RETURN v_value::numeric;
+exception
+    when others then
+        return null;
+END;
+$$;
+
+-- if we leak at least an integer in each subxact, then size of TopTransactionSize
+-- will be way beyond the 50k limit. If issue #3999 happens, then this will also take
+-- a long time, since for each row we will create a memory context that is not destroyed
+-- until the end of command.
+SELECT max(text2number('1234')), max(public.top_transaction_context_size()) > 50000 AS leaked
+FROM generate_series(1, 20000);
+
+-- Clean-up
+SET client_min_messages TO ERROR;
+DROP SCHEMA multi_subtransactions CASCADE;

src/test/regress/sql/multi_test_helpers_superuser.sql

@@ -48,3 +48,9 @@ CREATE OR REPLACE FUNCTION pg_catalog.partition_task_list_results(resultIdPrefix
                   targetShardIndex int)
     LANGUAGE C STRICT VOLATILE
     AS 'citus', $$partition_task_list_results$$;
+
+
+-- get size of TopTransactionContext
+CREATE OR REPLACE FUNCTION top_transaction_context_size() RETURNS BIGINT
+LANGUAGE C STRICT VOLATILE
+AS 'citus', $$top_transaction_context_size$$;