Create a stack of pl/python "execution contexts".

This replaces the former global variable PLy_curr_procedure, and provides
a place to stash per-call-level information.  In particular we create a
per-call-level scratch memory context.

For the moment, the scratch context is just used to avoid leaking memory
from datatype output function calls in PLyDict_FromTuple.  There probably
will be more use-cases in future.

Although this is a fix for a pre-existing memory leakage bug, it seems
sufficiently invasive to not want to back-patch; it feels better as part
of the major rearrangement of plpython code that we've already done as
part of 9.2.

Jan Urbański

src/pl/plpython/plpy_cursorobject.c

@@ -14,6 +14,7 @@
 #include "plpy_cursorobject.h"
 
 #include "plpy_elog.h"
+#include "plpy_main.h"
 #include "plpy_planobject.h"
 #include "plpy_procedure.h"
 #include "plpy_resultobject.h"
@@ -119,6 +120,7 @@ PLy_cursor_query(const char *query)
 
 	PG_TRY();
 	{
+		PLyExecutionContext	*exec_ctx = PLy_current_execution_context();
 		SPIPlanPtr	plan;
 		Portal		portal;
 
@@ -130,7 +132,7 @@ PLy_cursor_query(const char *query)
 				 SPI_result_code_string(SPI_result));
 
 		portal = SPI_cursor_open(NULL, plan, NULL, NULL,
-								 PLy_curr_procedure->fn_readonly);
+								 exec_ctx->curr_proc->fn_readonly);
 		SPI_freeplan(plan);
 
 		if (portal == NULL)
@@ -207,6 +209,7 @@ PLy_cursor_plan(PyObject *ob, PyObject *args)
 
 	PG_TRY();
 	{
+		PLyExecutionContext *exec_ctx = PLy_current_execution_context();
 		Portal		portal;
 		char	   *volatile nulls;
 		volatile int j;
@@ -253,7 +256,7 @@ PLy_cursor_plan(PyObject *ob, PyObject *args)
 		}
 
 		portal = SPI_cursor_open(NULL, plan->plan, plan->values, nulls,
-								 PLy_curr_procedure->fn_readonly);
+								 exec_ctx->curr_proc->fn_readonly);
 		if (portal == NULL)
 			elog(ERROR, "SPI_cursor_open() failed: %s",
 				 SPI_result_code_string(SPI_result));

src/pl/plpython/plpy_elog.c

@@ -12,6 +12,7 @@
 
 #include "plpy_elog.h"
 
+#include "plpy_main.h"
 #include "plpy_procedure.h"
 
 
@@ -255,6 +256,7 @@ PLy_traceback(char **xmsg, char **tbmsg, int *tb_depth)
 		/* The first frame always points at <module>, skip it. */
 		if (*tb_depth > 0)
 		{
+			PLyExecutionContext	*exec_ctx = PLy_current_execution_context();
 			char	   *proname;
 			char	   *fname;
 			char	   *line;
@@ -270,7 +272,7 @@ PLy_traceback(char **xmsg, char **tbmsg, int *tb_depth)
 			else
 				fname = PyString_AsString(name);
 
-			proname = PLy_procedure_name(PLy_curr_procedure);
+			proname = PLy_procedure_name(exec_ctx->curr_proc);
 			plain_filename = PyString_AsString(filename);
 			plain_lineno = PyInt_AsLong(lineno);
 
@@ -287,7 +289,7 @@ PLy_traceback(char **xmsg, char **tbmsg, int *tb_depth)
 			 * function code object was compiled with "<string>" as the
 			 * filename
 			 */
-			if (PLy_curr_procedure && plain_filename != NULL &&
+			if (exec_ctx->curr_proc && plain_filename != NULL &&
 				strcmp(plain_filename, "<string>") == 0)
 			{
 				/*
@@ -299,7 +301,7 @@ PLy_traceback(char **xmsg, char **tbmsg, int *tb_depth)
 				 * for.  But we do not go as far as traceback.py in reading
 				 * the source of imported modules.
 				 */
-				line = get_source_line(PLy_curr_procedure->src, plain_lineno);
+				line = get_source_line(exec_ctx->curr_proc->src, plain_lineno);
 				if (line)
 				{
 					appendStringInfo(&tbstr, "\n    %s", line);

src/pl/plpython/plpy_exec.c

@@ -455,7 +455,9 @@ PLy_function_delete_args(PLyProcedure *proc)
 static void
 plpython_return_error_callback(void *arg)
 {
-	if (PLy_curr_procedure)
+	PLyExecutionContext *exec_ctx = PLy_current_execution_context();
+
+	if (exec_ctx->curr_proc)
 		errcontext("while creating return value");
 }
 
@@ -781,7 +783,9 @@ PLy_modify_tuple(PLyProcedure *proc, PyObject *pltd, TriggerData *tdata,
 static void
 plpython_trigger_error_callback(void *arg)
 {
-	if (PLy_curr_procedure)
+	PLyExecutionContext *exec_ctx = PLy_current_execution_context();
+
+	if (exec_ctx->curr_proc)
 		errcontext("while modifying trigger row");
 }
 

src/pl/plpython/plpy_main.c

@@ -12,6 +12,7 @@
 #include "executor/spi.h"
 #include "miscadmin.h"
 #include "utils/guc.h"
+#include "utils/memutils.h"
 #include "utils/syscache.h"
 
 #include "plpython.h"
@@ -66,11 +67,17 @@ static void plpython_error_callback(void *arg);
 static void plpython_inline_error_callback(void *arg);
 static void PLy_init_interp(void);
 
+static PLyExecutionContext *PLy_push_execution_context(void);
+static void PLy_pop_execution_context(void);
+
 static const int plpython_python_version = PY_MAJOR_VERSION;
 
 /* initialize global variables */
 PyObject *PLy_interp_globals = NULL;
 
+/* this doesn't need to be global; use PLy_current_execution_context() */
+static PLyExecutionContext *PLy_execution_contexts = NULL;
+
 
 void
 _PG_init(void)
@@ -114,6 +121,8 @@ _PG_init(void)
 
 	explicit_subtransactions = NIL;
 
+	PLy_execution_contexts = NULL;
+
 	inited = true;
 }
 
@@ -179,13 +188,20 @@ Datum
 plpython_call_handler(PG_FUNCTION_ARGS)
 {
 	Datum		retval;
-	PLyProcedure *save_curr_proc;
+	PLyExecutionContext *exec_ctx;
 	ErrorContextCallback plerrcontext;
 
+	/* Note: SPI_finish() happens in plpy_exec.c, which is dubious design */
 	if (SPI_connect() != SPI_OK_CONNECT)
 		elog(ERROR, "SPI_connect failed");
 
-	save_curr_proc = PLy_curr_procedure;
+	/*
+	 * Push execution context onto stack.  It is important that this get
+	 * popped again, so avoid putting anything that could throw error between
+	 * here and the PG_TRY.  (plpython_error_callback expects the stack entry
+	 * to be there, so we have to make the context first.)
+	 */
+	exec_ctx = PLy_push_execution_context();
 
 	/*
 	 * Setup error traceback support for ereport()
@@ -203,20 +219,20 @@ plpython_call_handler(PG_FUNCTION_ARGS)
 			HeapTuple	trv;
 
 			proc = PLy_procedure_get(fcinfo->flinfo->fn_oid, true);
-			PLy_curr_procedure = proc;
+			exec_ctx->curr_proc = proc;
 			trv = PLy_exec_trigger(fcinfo, proc);
 			retval = PointerGetDatum(trv);
 		}
 		else
 		{
 			proc = PLy_procedure_get(fcinfo->flinfo->fn_oid, false);
-			PLy_curr_procedure = proc;
+			exec_ctx->curr_proc = proc;
 			retval = PLy_exec_function(fcinfo, proc);
 		}
 	}
 	PG_CATCH();
 	{
-		PLy_curr_procedure = save_curr_proc;
+		PLy_pop_execution_context();
 		PyErr_Clear();
 		PG_RE_THROW();
 	}
@@ -224,8 +240,8 @@ plpython_call_handler(PG_FUNCTION_ARGS)
 
 	/* Pop the error context stack */
 	error_context_stack = plerrcontext.previous;
-
+	/* ... and then the execution context */
-	PLy_curr_procedure = save_curr_proc;
+	PLy_pop_execution_context();
 
 	return retval;
 }
@@ -244,22 +260,14 @@ plpython_inline_handler(PG_FUNCTION_ARGS)
 	InlineCodeBlock *codeblock = (InlineCodeBlock *) DatumGetPointer(PG_GETARG_DATUM(0));
 	FunctionCallInfoData fake_fcinfo;
 	FmgrInfo	flinfo;
-	PLyProcedure *save_curr_proc;
 	PLyProcedure proc;
+	PLyExecutionContext *exec_ctx;
 	ErrorContextCallback plerrcontext;
 
+	/* Note: SPI_finish() happens in plpy_exec.c, which is dubious design */
 	if (SPI_connect() != SPI_OK_CONNECT)
 		elog(ERROR, "SPI_connect failed");
 
-	save_curr_proc = PLy_curr_procedure;
-
-	/*
-	 * Setup error traceback support for ereport()
-	 */
-	plerrcontext.callback = plpython_inline_error_callback;
-	plerrcontext.previous = error_context_stack;
-	error_context_stack = &plerrcontext;
-
 	MemSet(&fake_fcinfo, 0, sizeof(fake_fcinfo));
 	MemSet(&flinfo, 0, sizeof(flinfo));
 	fake_fcinfo.flinfo = &flinfo;
@@ -270,27 +278,44 @@ plpython_inline_handler(PG_FUNCTION_ARGS)
 	proc.pyname = PLy_strdup("__plpython_inline_block");
 	proc.result.out.d.typoid = VOIDOID;
 
+	/*
+	 * Push execution context onto stack.  It is important that this get
+	 * popped again, so avoid putting anything that could throw error between
+	 * here and the PG_TRY.  (plpython_inline_error_callback doesn't currently
+	 * need the stack entry, but for consistency with plpython_call_handler
+	 * we do it in this order.)
+	 */
+	exec_ctx = PLy_push_execution_context();
+
+	/*
+	 * Setup error traceback support for ereport()
+	 */
+	plerrcontext.callback = plpython_inline_error_callback;
+	plerrcontext.previous = error_context_stack;
+	error_context_stack = &plerrcontext;
+
 	PG_TRY();
 	{
 		PLy_procedure_compile(&proc, codeblock->source_text);
-		PLy_curr_procedure = &proc;
+		exec_ctx->curr_proc = &proc;
 		PLy_exec_function(&fake_fcinfo, &proc);
 	}
 	PG_CATCH();
 	{
+		PLy_pop_execution_context();
 		PLy_procedure_delete(&proc);
-		PLy_curr_procedure = save_curr_proc;
 		PyErr_Clear();
 		PG_RE_THROW();
 	}
 	PG_END_TRY();
 
-	PLy_procedure_delete(&proc);
-
 	/* Pop the error context stack */
 	error_context_stack = plerrcontext.previous;
+	/* ... and then the execution context */
+	PLy_pop_execution_context();
 
-	PLy_curr_procedure = save_curr_proc;
+	/* Now clean up the transient procedure we made */
+	PLy_procedure_delete(&proc);
 
 	PG_RETURN_VOID();
 }
@@ -313,9 +338,11 @@ static bool PLy_procedure_is_trigger(Form_pg_proc procStruct)
 static void
 plpython_error_callback(void *arg)
 {
-	if (PLy_curr_procedure)
+	PLyExecutionContext *exec_ctx = PLy_current_execution_context();
+
+	if (exec_ctx->curr_proc)
 		errcontext("PL/Python function \"%s\"",
-				   PLy_procedure_name(PLy_curr_procedure));
+				   PLy_procedure_name(exec_ctx->curr_proc));
 }
 
 static void
@@ -323,3 +350,42 @@ plpython_inline_error_callback(void *arg)
 {
 	errcontext("PL/Python anonymous code block");
 }
+
+PLyExecutionContext *
+PLy_current_execution_context(void)
+{
+	if (PLy_execution_contexts == NULL)
+		elog(ERROR, "no Python function is currently executing");
+
+	return PLy_execution_contexts;
+}
+
+static PLyExecutionContext *
+PLy_push_execution_context(void)
+{
+	PLyExecutionContext	*context = PLy_malloc(sizeof(PLyExecutionContext));
+
+	context->curr_proc = NULL;
+	context->scratch_ctx = AllocSetContextCreate(TopTransactionContext,
+												 "PL/Python scratch context",
+												 ALLOCSET_DEFAULT_MINSIZE,
+												 ALLOCSET_DEFAULT_INITSIZE,
+												 ALLOCSET_DEFAULT_MAXSIZE);
+	context->next = PLy_execution_contexts;
+	PLy_execution_contexts = context;
+	return context;
+}
+
+static void
+PLy_pop_execution_context(void)
+{
+	PLyExecutionContext	*context = PLy_execution_contexts;
+
+	if (context == NULL)
+		elog(ERROR, "no Python function is currently executing");
+
+	PLy_execution_contexts = context->next;
+
+	MemoryContextDelete(context->scratch_ctx);
+	PLy_free(context);
+}

src/pl/plpython/plpy_main.h

@@ -10,4 +10,19 @@
 /* the interpreter's globals dict */
 extern PyObject *PLy_interp_globals;
 
+/*
+ * A stack of PL/Python execution contexts. Each time user-defined Python code
+ * is called, an execution context is created and put on the stack. After the
+ * Python code returns, the context is destroyed.
+ */
+typedef struct PLyExecutionContext
+{
+	PLyProcedure	*curr_proc;		/* the currently executing procedure */
+	MemoryContext	scratch_ctx;	/* a context for things like type I/O */
+	struct PLyExecutionContext *next;	/* previous stack level */
+} PLyExecutionContext;
+
+/* Get the current execution context */
+extern PLyExecutionContext *PLy_current_execution_context(void);
+
 #endif	/* PLPY_MAIN_H */

src/pl/plpython/plpy_procedure.c

@@ -22,9 +22,6 @@
 #include "plpy_main.h"
 
 
-PLyProcedure *PLy_curr_procedure = NULL;
-
-
 static HTAB *PLy_procedure_cache = NULL;
 static HTAB *PLy_trigger_cache = NULL;
 

src/pl/plpython/plpy_procedure.h

@@ -45,8 +45,4 @@ extern PLyProcedure *PLy_procedure_get(Oid fn_oid, bool is_trigger);
 extern void PLy_procedure_compile(PLyProcedure *proc, const char *src);
 extern void PLy_procedure_delete(PLyProcedure *proc);
 
-
-/* currently active plpython function */
-extern PLyProcedure *PLy_curr_procedure;
-
 #endif	/* PLPY_PROCEDURE_H */

src/pl/plpython/plpy_spi.c

@@ -18,6 +18,7 @@
 #include "plpy_spi.h"
 
 #include "plpy_elog.h"
+#include "plpy_main.h"
 #include "plpy_planobject.h"
 #include "plpy_plpymodule.h"
 #include "plpy_procedure.h"
@@ -236,6 +237,7 @@ PLy_spi_execute_plan(PyObject *ob, PyObject *list, long limit)
 
 	PG_TRY();
 	{
+		PLyExecutionContext *exec_ctx = PLy_current_execution_context();
 		char	   *volatile nulls;
 		volatile int j;
 
@@ -281,7 +283,7 @@ PLy_spi_execute_plan(PyObject *ob, PyObject *list, long limit)
 		}
 
 		rv = SPI_execute_plan(plan->plan, plan->values, nulls,
-							  PLy_curr_procedure->fn_readonly, limit);
+							  exec_ctx->curr_proc->fn_readonly, limit);
 		ret = PLy_spi_execute_fetch_result(SPI_tuptable, SPI_processed, rv);
 
 		if (nargs > 0)
@@ -347,8 +349,10 @@ PLy_spi_execute_query(char *query, long limit)
 
 	PG_TRY();
 	{
+		PLyExecutionContext	*exec_ctx = PLy_current_execution_context();
+
 		pg_verifymbstr(query, strlen(query), false);
-		rv = SPI_execute(query, PLy_curr_procedure->fn_readonly, limit);
+		rv = SPI_execute(query, exec_ctx->curr_proc->fn_readonly, limit);
 		ret = PLy_spi_execute_fetch_result(SPI_tuptable, SPI_processed, rv);
 
 		PLy_spi_subtransaction_commit(oldcontext, oldowner);

src/pl/plpython/plpy_typeio.c

@@ -23,6 +23,7 @@
 #include "plpy_typeio.h"
 
 #include "plpy_elog.h"
+#include "plpy_main.h"
 
 
 /* I/O function caching */
@@ -258,10 +259,15 @@ PLy_output_record_funcs(PLyTypeInfo *arg, TupleDesc desc)
 	Assert(arg->is_rowtype == 1);
 }
 
+/*
+ * Transform a tuple into a Python dict object.
+ */
 PyObject *
 PLyDict_FromTuple(PLyTypeInfo *info, HeapTuple tuple, TupleDesc desc)
 {
 	PyObject   *volatile dict;
+	PLyExecutionContext *exec_ctx = PLy_current_execution_context();
+	MemoryContext oldcontext = CurrentMemoryContext;
 	int			i;
 
 	if (info->is_rowtype != 1)
@@ -273,6 +279,11 @@ PLyDict_FromTuple(PLyTypeInfo *info, HeapTuple tuple, TupleDesc desc)
 
 	PG_TRY();
 	{
+		/*
+		 * Do the work in the scratch context to avoid leaking memory from
+		 * the datatype output function calls.
+		 */
+		MemoryContextSwitchTo(exec_ctx->scratch_ctx);
 		for (i = 0; i < info->in.r.natts; i++)
 		{
 			char	   *key;
@@ -295,9 +306,12 @@ PLyDict_FromTuple(PLyTypeInfo *info, HeapTuple tuple, TupleDesc desc)
 				Py_DECREF(value);
 			}
 		}
+		MemoryContextSwitchTo(oldcontext);
+		MemoryContextReset(exec_ctx->scratch_ctx);
 	}
 	PG_CATCH();
 	{
+		MemoryContextSwitchTo(oldcontext);
 		Py_DECREF(dict);
 		PG_RE_THROW();
 	}