From 3651fc64ee27add6c82fc1097ab27ddb49b27ce3 Mon Sep 17 00:00:00 2001 From: Hadi Moshayedi Date: Wed, 8 Jul 2020 19:42:39 -0700 Subject: [PATCH] Fix Subtransaction memory leak --- .../distributed/commands/variableset.c | 3 +- src/backend/distributed/test/xact_stats.c | 56 +++++++++++++++++++ .../transaction/transaction_management.c | 45 +++++++++++---- .../expected/multi_subtransactions.out | 33 +++++++++-- .../expected/multi_test_helpers_superuser.out | 4 ++ .../regress/sql/multi_subtransactions.sql | 29 +++++++++- .../sql/multi_test_helpers_superuser.sql | 6 ++ 7 files changed, 158 insertions(+), 18 deletions(-) create mode 100644 src/backend/distributed/test/xact_stats.c diff --git a/src/backend/distributed/commands/variableset.c b/src/backend/distributed/commands/variableset.c index b0fc125c9..9f93e7c65 100644 --- a/src/backend/distributed/commands/variableset.c +++ b/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); } diff --git a/src/backend/distributed/test/xact_stats.c b/src/backend/distributed/test/xact_stats.c new file mode 100644 index 000000000..b157ea27a --- /dev/null +++ b/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 +#include + +#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; +} diff --git a/src/backend/distributed/transaction/transaction_management.c b/src/backend/distributed/transaction/transaction_management.c index cfd35fbe6..0e39ca905 100644 --- a/src/backend/distributed/transaction/transaction_management.c +++ b/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); } diff --git a/src/test/regress/expected/multi_subtransactions.out b/src/test/regress/expected/multi_subtransactions.out index aaa8f2ef5..069ad3524 100644 --- a/src/test/regress/expected/multi_subtransactions.out +++ b/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; diff --git a/src/test/regress/expected/multi_test_helpers_superuser.out b/src/test/regress/expected/multi_test_helpers_superuser.out index d80f310b9..b631814f8 100644 --- a/src/test/regress/expected/multi_test_helpers_superuser.out +++ b/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$$; diff --git a/src/test/regress/sql/multi_subtransactions.sql b/src/test/regress/sql/multi_subtransactions.sql index 5d51272c1..5ec750644 100644 --- a/src/test/regress/sql/multi_subtransactions.sql +++ b/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; diff --git a/src/test/regress/sql/multi_test_helpers_superuser.sql b/src/test/regress/sql/multi_test_helpers_superuser.sql index 748e9ef89..aa7b3ee66 100644 --- a/src/test/regress/sql/multi_test_helpers_superuser.sql +++ b/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$$;