Prevent leakage of cached plans and execution trees in plpgsql DO blocks.

plpgsql likes to cache query plans and simple-expression execution state
trees across calls.  This is a considerable win for multiple executions
of the same function.  However, it's useless for DO blocks, since by
definition those are executed only once and discarded.  Nonetheless,
we were allowing a DO block's expression execution trees to survive
until end of transaction, resulting in a significant intra-transaction
memory leak, as reported by Yeb Havinga.  Worse, if the DO block exited
with an error, the compiled form of the block's code was leaked till
end of session --- along with subsidiary plancache entries.

To fix, make DO blocks keep their expression execution trees in a private
EState that's deleted at exit from the block, and add a PG_TRY block
to plpgsql_inline_handler to make sure that memory cleanup happens
even on error exits.  Also add a regression test covering error handling
in a DO block, because my first try at this broke that.  (The test is
not meant to prove that we don't leak memory anymore, though it could
be used for that with a much larger loop count.)

Ideally we'd back-patch this into all versions supporting DO blocks;
but the patch needs to add a field to struct PLpgSQL_execstate, and that
would break ABI compatibility for third-party plugins such as the plpgsql
debugger.  Given the small number of complaints so far, fixing this in
HEAD only seems like an acceptable choice.

5 files changed, 150 insertions, 22 deletions

diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
index bc31fe90850..f5f1892e69d 100644
--- a/src/pl/plpgsql/src/pl_exec.c
+++ b/src/pl/plpgsql/src/pl_exec.c
@@ -66,6 +66,15 @@ typedef struct
  * so that we don't have to re-prepare simple expressions on each trip through
  * a function.	(We assume the case to optimize is many repetitions of a
  * function within a transaction.)
+ *
+ * However, there's no value in trying to amortize simple expression setup
+ * across multiple executions of a DO block (inline code block), since there
+ * can never be any.  If we use the shared EState for a DO block, the expr
+ * state trees are effectively leaked till end of transaction, and that can
+ * add up if the user keeps on submitting DO blocks.  Therefore, each DO block
+ * has its own simple-expression EState, which is cleaned up at exit from
+ * plpgsql_inline_handler().  DO blocks still use the simple_econtext_stack,
+ * though, so that subxact abort cleanup does the right thing.
  */
 typedef struct SimpleEcontextStackEntry
 {
@@ -74,7 +83,7 @@ typedef struct SimpleEcontextStackEntry
 	struct SimpleEcontextStackEntry *next;		/* next stack entry up */
 } SimpleEcontextStackEntry;
 
-static EState *simple_eval_estate = NULL;
+static EState *shared_simple_eval_estate = NULL;
 static SimpleEcontextStackEntry *simple_econtext_stack = NULL;
 
 /************************************************************
@@ -136,7 +145,8 @@ static int exec_stmt_dynfors(PLpgSQL_execstate *estate,
 
 static void plpgsql_estate_setup(PLpgSQL_execstate *estate,
 					 PLpgSQL_function *func,
-					 ReturnSetInfo *rsi);
+					 ReturnSetInfo *rsi,
+					 EState *simple_eval_estate);
 static void exec_eval_cleanup(PLpgSQL_execstate *estate);
 
 static void exec_prepare_plan(PLpgSQL_execstate *estate,
@@ -230,10 +240,17 @@ static char *format_preparedparamsdata(PLpgSQL_execstate *estate,
 /* ----------
  * plpgsql_exec_function	Called by the call handler for
  *				function execution.
+ *
+ * This is also used to execute inline code blocks (DO blocks).  The only
+ * difference that this code is aware of is that for a DO block, we want
+ * to use a private simple_eval_estate, which is created and passed in by
+ * the caller.  For regular functions, pass NULL, which implies using
+ * shared_simple_eval_estate.
  * ----------
  */
 Datum
-plpgsql_exec_function(PLpgSQL_function *func, FunctionCallInfo fcinfo)
+plpgsql_exec_function(PLpgSQL_function *func, FunctionCallInfo fcinfo,
+					  EState *simple_eval_estate)
 {
 	PLpgSQL_execstate estate;
 	ErrorContextCallback plerrcontext;
@@ -243,7 +260,8 @@ plpgsql_exec_function(PLpgSQL_function *func, FunctionCallInfo fcinfo)
 	/*
 	 * Setup the execution state
 	 */
-	plpgsql_estate_setup(&estate, func, (ReturnSetInfo *) fcinfo->resultinfo);
+	plpgsql_estate_setup(&estate, func, (ReturnSetInfo *) fcinfo->resultinfo,
+						 simple_eval_estate);
 
 	/*
 	 * Setup error traceback support for ereport()
@@ -503,7 +521,7 @@ plpgsql_exec_trigger(PLpgSQL_function *func,
 	/*
 	 * Setup the execution state
 	 */
-	plpgsql_estate_setup(&estate, func, NULL);
+	plpgsql_estate_setup(&estate, func, NULL, NULL);
 
 	/*
 	 * Setup error traceback support for ereport()
@@ -782,7 +800,7 @@ plpgsql_exec_event_trigger(PLpgSQL_function *func, EventTriggerData *trigdata)
 	/*
 	 * Setup the execution state
 	 */
-	plpgsql_estate_setup(&estate, func, NULL);
+	plpgsql_estate_setup(&estate, func, NULL, NULL);
 
 	/*
 	 * Setup error traceback support for ereport()
@@ -3087,7 +3105,8 @@ exec_stmt_raise(PLpgSQL_execstate *estate, PLpgSQL_stmt_raise *stmt)
 static void
 plpgsql_estate_setup(PLpgSQL_execstate *estate,
 					 PLpgSQL_function *func,
-					 ReturnSetInfo *rsi)
+					 ReturnSetInfo *rsi,
+					 EState *simple_eval_estate)
 {
 	/* this link will be restored at exit from plpgsql_call_handler */
 	func->cur_estate = estate;
@@ -3126,6 +3145,12 @@ plpgsql_estate_setup(PLpgSQL_execstate *estate,
 	estate->datums = palloc(sizeof(PLpgSQL_datum *) * estate->ndatums);
 	/* caller is expected to fill the datums array */
 
+	/* set up for use of appropriate simple-expression EState */
+	if (simple_eval_estate)
+		estate->simple_eval_estate = simple_eval_estate;
+	else
+		estate->simple_eval_estate = shared_simple_eval_estate;
+
 	estate->eval_tuptable = NULL;
 	estate->eval_processed = 0;
 	estate->eval_lastoid = InvalidOid;
@@ -5148,7 +5173,7 @@ exec_eval_simple_expr(PLpgSQL_execstate *estate,
 	 */
 	if (expr->expr_simple_lxid != curlxid)
 	{
-		oldcontext = MemoryContextSwitchTo(simple_eval_estate->es_query_cxt);
+		oldcontext = MemoryContextSwitchTo(estate->simple_eval_estate->es_query_cxt);
 		expr->expr_simple_state = ExecInitExpr(expr->expr_simple_expr, NULL);
 		expr->expr_simple_in_use = false;
 		expr->expr_simple_lxid = curlxid;
@@ -6190,8 +6215,8 @@ exec_set_found(PLpgSQL_execstate *estate, bool state)
 /*
  * plpgsql_create_econtext --- create an eval_econtext for the current function
  *
- * We may need to create a new simple_eval_estate too, if there's not one
- * already for the current transaction.  The EState will be cleaned up at
+ * We may need to create a new shared_simple_eval_estate too, if there's not
+ * one already for the current transaction.  The EState will be cleaned up at
  * transaction end.
  */
 static void
@@ -6203,20 +6228,25 @@ plpgsql_create_econtext(PLpgSQL_execstate *estate)
 	 * Create an EState for evaluation of simple expressions, if there's not
 	 * one already in the current transaction.	The EState is made a child of
 	 * TopTransactionContext so it will have the right lifespan.
+	 *
+	 * Note that this path is never taken when executing a DO block; the
+	 * required EState was already made by plpgsql_inline_handler.
 	 */
-	if (simple_eval_estate == NULL)
+	if (estate->simple_eval_estate == NULL)
 	{
 		MemoryContext oldcontext;
 
+		Assert(shared_simple_eval_estate == NULL);
 		oldcontext = MemoryContextSwitchTo(TopTransactionContext);
-		simple_eval_estate = CreateExecutorState();
+		shared_simple_eval_estate = CreateExecutorState();
+		estate->simple_eval_estate = shared_simple_eval_estate;
 		MemoryContextSwitchTo(oldcontext);
 	}
 
 	/*
 	 * Create a child econtext for the current function.
 	 */
-	estate->eval_econtext = CreateExprContext(simple_eval_estate);
+	estate->eval_econtext = CreateExprContext(estate->simple_eval_estate);
 
 	/*
 	 * Make a stack entry so we can clean up the econtext at subxact end.
@@ -6275,14 +6305,14 @@ plpgsql_xact_cb(XactEvent event, void *arg)
 		/* Shouldn't be any econtext stack entries left at commit */
 		Assert(simple_econtext_stack == NULL);
 
-		if (simple_eval_estate)
-			FreeExecutorState(simple_eval_estate);
-		simple_eval_estate = NULL;
+		if (shared_simple_eval_estate)
+			FreeExecutorState(shared_simple_eval_estate);
+		shared_simple_eval_estate = NULL;
 	}
 	else if (event == XACT_EVENT_ABORT)
 	{
 		simple_econtext_stack = NULL;
-		simple_eval_estate = NULL;
+		shared_simple_eval_estate = NULL;
 	}
 }
 
@@ -6291,8 +6321,7 @@ plpgsql_xact_cb(XactEvent event, void *arg)
  *
  * Make sure any simple-expression econtexts created in the current
  * subtransaction get cleaned up.  We have to do this explicitly because
- * no other code knows which child econtexts of simple_eval_estate belong
- * to which level of subxact.
+ * no other code knows which econtexts belong to which level of subxact.
  */
 void
 plpgsql_subxact_cb(SubXactEvent event, SubTransactionId mySubid,
diff --git a/src/pl/plpgsql/src/pl_handler.c b/src/pl/plpgsql/src/pl_handler.c
index 912422cd2eb..5bfe3c3d8ef 100644
--- a/src/pl/plpgsql/src/pl_handler.c
+++ b/src/pl/plpgsql/src/pl_handler.c
@@ -136,7 +136,7 @@ plpgsql_call_handler(PG_FUNCTION_ARGS)
 			retval = (Datum) 0;
 		}
 		else
-			retval = plpgsql_exec_function(func, fcinfo);
+			retval = plpgsql_exec_function(func, fcinfo, NULL);
 	}
 	PG_CATCH();
 	{
@@ -175,6 +175,7 @@ plpgsql_inline_handler(PG_FUNCTION_ARGS)
 	PLpgSQL_function *func;
 	FunctionCallInfoData fake_fcinfo;
 	FmgrInfo	flinfo;
+	EState	   *simple_eval_estate;
 	Datum		retval;
 	int			rc;
 
@@ -203,7 +204,51 @@ plpgsql_inline_handler(PG_FUNCTION_ARGS)
 	flinfo.fn_oid = InvalidOid;
 	flinfo.fn_mcxt = CurrentMemoryContext;
 
-	retval = plpgsql_exec_function(func, &fake_fcinfo);
+	/* Create a private EState for simple-expression execution */
+	simple_eval_estate = CreateExecutorState();
+
+	/* And run the function */
+	PG_TRY();
+	{
+		retval = plpgsql_exec_function(func, &fake_fcinfo, simple_eval_estate);
+	}
+	PG_CATCH();
+	{
+		/*
+		 * We need to clean up what would otherwise be long-lived resources
+		 * accumulated by the failed DO block, principally cached plans for
+		 * statements (which can be flushed with plpgsql_free_function_memory)
+		 * and execution trees for simple expressions, which are in the
+		 * private EState.
+		 *
+		 * Before releasing the private EState, we must clean up any
+		 * simple_econtext_stack entries pointing into it, which we can do by
+		 * invoking the subxact callback.  (It will be called again later if
+		 * some outer control level does a subtransaction abort, but no harm
+		 * is done.)  We cheat a bit knowing that plpgsql_subxact_cb does not
+		 * pay attention to its parentSubid argument.
+		 */
+		plpgsql_subxact_cb(SUBXACT_EVENT_ABORT_SUB,
+						   GetCurrentSubTransactionId(),
+						   0, NULL);
+
+		/* Clean up the private EState */
+		FreeExecutorState(simple_eval_estate);
+
+		/* Function should now have no remaining use-counts ... */
+		func->use_count--;
+		Assert(func->use_count == 0);
+
+		/* ... so we can free subsidiary storage */
+		plpgsql_free_function_memory(func);
+
+		/* And propagate the error */
+		PG_RE_THROW();
+	}
+	PG_END_TRY();
+
+	/* Clean up the private EState */
+	FreeExecutorState(simple_eval_estate);
 
 	/* Function should now have no remaining use-counts ... */
 	func->use_count--;
diff --git a/src/pl/plpgsql/src/plpgsql.h b/src/pl/plpgsql/src/plpgsql.h
index 9cb4f533346..3afcdf324cc 100644
--- a/src/pl/plpgsql/src/plpgsql.h
+++ b/src/pl/plpgsql/src/plpgsql.h
@@ -773,10 +773,14 @@ typedef struct PLpgSQL_execstate
 	ResourceOwner tuple_store_owner;
 	ReturnSetInfo *rsi;
 
+	/* the datums representing the function's local variables */
 	int			found_varno;
 	int			ndatums;
 	PLpgSQL_datum **datums;
 
+	/* EState to use for "simple" expression evaluation */
+	EState	   *simple_eval_estate;
+
 	/* temporary state for results from evaluation of query or expr */
 	SPITupleTable *eval_tuptable;
 	uint32		eval_processed;
@@ -943,7 +947,8 @@ extern Datum plpgsql_validator(PG_FUNCTION_ARGS);
  * ----------
  */
 extern Datum plpgsql_exec_function(PLpgSQL_function *func,
-					  FunctionCallInfo fcinfo);
+					  FunctionCallInfo fcinfo,
+					  EState *simple_eval_estate);
 extern HeapTuple plpgsql_exec_trigger(PLpgSQL_function *func,
 					 TriggerData *trigdata);
 extern void plpgsql_exec_event_trigger(PLpgSQL_function *func,
diff --git a/src/test/regress/expected/plpgsql.out b/src/test/regress/expected/plpgsql.out
index 22c16c4b3e0..0405ef47dd8 100644
--- a/src/test/regress/expected/plpgsql.out
+++ b/src/test/regress/expected/plpgsql.out
@@ -4651,6 +4651,35 @@ LINE 1: SELECT rtrim(roomno) AS roomno, foo FROM Room ORDER BY roomn...
                                         ^
 QUERY:  SELECT rtrim(roomno) AS roomno, foo FROM Room ORDER BY roomno
 CONTEXT:  PL/pgSQL function inline_code_block line 4 at FOR over SELECT rows
+-- Check handling of errors thrown from/into anonymous code blocks.
+do $outer$
+begin
+  for i in 1..10 loop
+   begin
+    execute $ex$
+      do $$
+      declare x int = 0;
+      begin
+        x := 1 / x;
+      end;
+      $$;
+    $ex$;
+  exception when division_by_zero then
+    raise notice 'caught division by zero';
+  end;
+  end loop;
+end;
+$outer$;
+NOTICE:  caught division by zero
+NOTICE:  caught division by zero
+NOTICE:  caught division by zero
+NOTICE:  caught division by zero
+NOTICE:  caught division by zero
+NOTICE:  caught division by zero
+NOTICE:  caught division by zero
+NOTICE:  caught division by zero
+NOTICE:  caught division by zero
+NOTICE:  caught division by zero
 -- Check variable scoping -- a var is not available in its own or prior
 -- default expressions.
 create function scope_test() returns int as $$
diff --git a/src/test/regress/sql/plpgsql.sql b/src/test/regress/sql/plpgsql.sql
index a685fa246e3..d01011a85a7 100644
--- a/src/test/regress/sql/plpgsql.sql
+++ b/src/test/regress/sql/plpgsql.sql
@@ -3751,6 +3751,26 @@ BEGIN
     END LOOP;
 END$$;
 
+-- Check handling of errors thrown from/into anonymous code blocks.
+do $outer$
+begin
+  for i in 1..10 loop
+   begin
+    execute $ex$
+      do $$
+      declare x int = 0;
+      begin
+        x := 1 / x;
+      end;
+      $$;
+    $ex$;
+  exception when division_by_zero then
+    raise notice 'caught division by zero';
+  end;
+  end loop;
+end;
+$outer$;
+
 -- Check variable scoping -- a var is not available in its own or prior
 -- default expressions.