Previous fix for temporary file management broke returning a set from

PL/pgSQL function within an exception handler. Make sure we use the right
resource owner when we create the tuplestore to hold returned tuples.

Simplify tuplestore API so that the caller doesn't need to be in the right
memory context when calling tuplestore_put* functions. tuplestore.c
automatically switches to the memory context used when the tuplestore was
created. Tuplesort was already modified like this earlier. This patch also
removes the now useless MemoryContextSwitch calls from callers.

Report by Aleksei on pgsql-bugs on Dec 22 2009. Backpatch to 8.1, like
the previous patch that broke this.

src/backend/commands/prepare.c

@@ -10,7 +10,7 @@
  * Copyright (c) 2002-2009, PostgreSQL Global Development Group
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/commands/prepare.c,v 1.97 2009/06/11 14:48:56 momjian Exp $
+ *	  $PostgreSQL: pgsql/src/backend/commands/prepare.c,v 1.97.2.1 2009/12/29 17:41:09 heikki Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -781,6 +781,9 @@ pg_prepared_statement(PG_FUNCTION_ARGS)
 		tuplestore_begin_heap(rsinfo->allowedModes & SFRM_Materialize_Random,
 							  false, work_mem);
 
+	/* generate junk in short-term context */
+	MemoryContextSwitchTo(oldcontext);
+
 	/* hash table might be uninitialized */
 	if (prepared_queries)
 	{
@@ -793,9 +796,6 @@ pg_prepared_statement(PG_FUNCTION_ARGS)
 			Datum		values[5];
 			bool		nulls[5];
 
-			/* generate junk in short-term context */
-			MemoryContextSwitchTo(oldcontext);
-
 			MemSet(nulls, 0, sizeof(nulls));
 
 			values[0] = CStringGetTextDatum(prep_stmt->stmt_name);
@@ -805,8 +805,6 @@ pg_prepared_statement(PG_FUNCTION_ARGS)
 										  prep_stmt->plansource->num_params);
 			values[4] = BoolGetDatum(prep_stmt->from_sql);
 
-			/* switch to appropriate context while storing the tuple */
-			MemoryContextSwitchTo(per_query_ctx);
 			tuplestore_putvalues(tupstore, tupdesc, values, nulls);
 		}
 	}
@@ -814,8 +812,6 @@ pg_prepared_statement(PG_FUNCTION_ARGS)
 	/* clean up and return the tuplestore */
 	tuplestore_donestoring(tupstore);
 
-	MemoryContextSwitchTo(oldcontext);
-
 	rsinfo->returnMode = SFRM_Materialize;
 	rsinfo->setResult = tupstore;
 	rsinfo->setDesc = tupdesc;

src/backend/executor/execQual.c

@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/executor/execQual.c,v 1.250.2.1 2009/12/14 02:16:03 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/backend/executor/execQual.c,v 1.250.2.2 2009/12/29 17:41:09 heikki Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -2000,15 +2000,10 @@ ExecMakeTableFunctionResult(ExprState *funcexpr,
 				tmptup.t_len = HeapTupleHeaderGetDatumLength(td);
 				tmptup.t_data = td;
 
-				oldcontext = MemoryContextSwitchTo(econtext->ecxt_per_query_memory);
 				tuplestore_puttuple(tupstore, &tmptup);
 			}
 			else
-			{
-				oldcontext = MemoryContextSwitchTo(econtext->ecxt_per_query_memory);
 				tuplestore_putvalues(tupstore, tupdesc, &result, &fcinfo.isnull);
-			}
-			MemoryContextSwitchTo(oldcontext);
 
 			/*
 			 * Are we done?

src/backend/executor/functions.c

@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/executor/functions.c,v 1.135.2.1 2009/12/14 02:16:03 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/backend/executor/functions.c,v 1.135.2.2 2009/12/29 17:41:09 heikki Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -1400,15 +1400,12 @@ static void
 sqlfunction_receive(TupleTableSlot *slot, DestReceiver *self)
 {
 	DR_sqlfunction *myState = (DR_sqlfunction *) self;
-	MemoryContext oldcxt;
 
 	/* Filter tuple as needed */
 	slot = ExecFilterJunk(myState->filter, slot);
 
 	/* Store the filtered tuple into the tuplestore */
-	oldcxt = MemoryContextSwitchTo(myState->cxt);
 	tuplestore_puttupleslot(myState->tstore, slot);
-	MemoryContextSwitchTo(oldcxt);
 }
 
 /*

src/backend/executor/nodeWindowAgg.c

@@ -27,7 +27,7 @@
  * Portions Copyright (c) 1994, Regents of the University of California
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/executor/nodeWindowAgg.c,v 1.6 2009/06/20 18:45:28 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/backend/executor/nodeWindowAgg.c,v 1.6.2.1 2009/12/29 17:41:09 heikki Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -723,7 +723,7 @@ spool_tuples(WindowAggState *winstate, int64 pos)
 
 	outerPlan = outerPlanState(winstate);
 
-	/* Must be in query context to call outerplan or touch tuplestore */
+	/* Must be in query context to call outerplan */
 	oldcontext = MemoryContextSwitchTo(winstate->ss.ps.ps_ExprContext->ecxt_per_query_memory);
 
 	while (winstate->spooled_rows <= pos || pos == -1)

src/backend/executor/tstoreReceiver.c

@@ -13,7 +13,7 @@
  * Portions Copyright (c) 1994, Regents of the University of California
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/executor/tstoreReceiver.c,v 1.23 2009/06/11 14:48:57 momjian Exp $
+ *	  $PostgreSQL: pgsql/src/backend/executor/tstoreReceiver.c,v 1.23.2.1 2009/12/29 17:41:09 heikki Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -94,11 +94,8 @@ static void
 tstoreReceiveSlot_notoast(TupleTableSlot *slot, DestReceiver *self)
 {
 	TStoreState *myState = (TStoreState *) self;
-	MemoryContext oldcxt = MemoryContextSwitchTo(myState->cxt);
 
 	tuplestore_puttupleslot(myState->tstore, slot);
-
-	MemoryContextSwitchTo(oldcxt);
 }
 
 /*

src/backend/utils/mmgr/portalmem.c

@@ -12,7 +12,7 @@
  * Portions Copyright (c) 1994, Regents of the University of California
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/utils/mmgr/portalmem.c,v 1.113 2009/01/01 17:23:53 momjian Exp $
+ *	  $PostgreSQL: pgsql/src/backend/utils/mmgr/portalmem.c,v 1.113.2.1 2009/12/29 17:41:09 heikki Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -923,6 +923,9 @@ pg_cursor(PG_FUNCTION_ARGS)
 		tuplestore_begin_heap(rsinfo->allowedModes & SFRM_Materialize_Random,
 							  false, work_mem);
 
+	/* generate junk in short-term context */
+	MemoryContextSwitchTo(oldcontext);
+
 	hash_seq_init(&hash_seq, PortalHashTable);
 	while ((hentry = hash_seq_search(&hash_seq)) != NULL)
 	{
@@ -934,9 +937,6 @@ pg_cursor(PG_FUNCTION_ARGS)
 		if (!portal->visible)
 			continue;
 
-		/* generate junk in short-term context */
-		MemoryContextSwitchTo(oldcontext);
-
 		MemSet(nulls, 0, sizeof(nulls));
 
 		values[0] = CStringGetTextDatum(portal->name);
@@ -946,16 +946,12 @@ pg_cursor(PG_FUNCTION_ARGS)
 		values[4] = BoolGetDatum(portal->cursorOptions & CURSOR_OPT_SCROLL);
 		values[5] = TimestampTzGetDatum(portal->creation_time);
 
-		/* switch to appropriate context while storing the tuple */
-		MemoryContextSwitchTo(per_query_ctx);
 		tuplestore_putvalues(tupstore, tupdesc, values, nulls);
 	}
 
 	/* clean up and return the tuplestore */
 	tuplestore_donestoring(tupstore);
 
-	MemoryContextSwitchTo(oldcontext);
-
 	rsinfo->returnMode = SFRM_Materialize;
 	rsinfo->setResult = tupstore;
 	rsinfo->setDesc = tupdesc;

src/backend/utils/sort/tuplestore.c

@@ -47,7 +47,7 @@
  * Portions Copyright (c) 1994, Regents of the University of California
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/utils/sort/tuplestore.c,v 1.48 2009/06/11 14:49:06 momjian Exp $
+ *	  $PostgreSQL: pgsql/src/backend/utils/sort/tuplestore.c,v 1.48.2.1 2009/12/29 17:41:09 heikki Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -58,6 +58,7 @@
 #include "executor/executor.h"
 #include "storage/buffile.h"
 #include "utils/memutils.h"
+#include "utils/resowner.h"
 #include "utils/tuplestore.h"
 
 
@@ -105,6 +106,8 @@ struct Tuplestorestate
 	bool		truncated;		/* tuplestore_trim has removed tuples? */
 	long		availMem;		/* remaining memory available, in bytes */
 	BufFile    *myfile;			/* underlying file, or NULL if none */
+	MemoryContext context;		/* memory context for holding tuples */
+	ResourceOwner resowner;		/* resowner for holding temp files */
 
 	/*
 	 * These function pointers decouple the routines that must know what kind
@@ -246,6 +249,8 @@ tuplestore_begin_common(int eflags, bool interXact, int maxKBytes)
 	state->truncated = false;
 	state->availMem = maxKBytes * 1024L;
 	state->myfile = NULL;
+	state->context = CurrentMemoryContext;
+	state->resowner = CurrentResourceOwner;
 
 	state->memtupcount = 0;
 	state->memtupsize = 1024;	/* initial guess */
@@ -278,9 +283,9 @@ tuplestore_begin_common(int eflags, bool interXact, int maxKBytes)
  *
  * interXact: if true, the files used for on-disk storage persist beyond the
  * end of the current transaction.	NOTE: It's the caller's responsibility to
- * create such a tuplestore in a memory context that will also survive
- * transaction boundaries, and to ensure the tuplestore is closed when it's
- * no longer wanted.
+ * create such a tuplestore in a memory context and resource owner that will
+ * also survive transaction boundaries, and to ensure the tuplestore is closed
+ * when it's no longer wanted.
  *
  * maxKBytes: how much data to store in memory (any data beyond this
  * amount is paged to disk).  When in doubt, use work_mem.
@@ -533,6 +538,7 @@ tuplestore_puttupleslot(Tuplestorestate *state,
 						TupleTableSlot *slot)
 {
 	MinimalTuple tuple;
+	MemoryContext oldcxt = MemoryContextSwitchTo(state->context);
 
 	/*
 	 * Form a MinimalTuple in working memory
@@ -541,6 +547,8 @@ tuplestore_puttupleslot(Tuplestorestate *state,
 	USEMEM(state, GetMemoryChunkSpace(tuple));
 
 	tuplestore_puttuple_common(state, (void *) tuple);
+
+	MemoryContextSwitchTo(oldcxt);
 }
 
 /*
@@ -550,12 +558,16 @@ tuplestore_puttupleslot(Tuplestorestate *state,
 void
 tuplestore_puttuple(Tuplestorestate *state, HeapTuple tuple)
 {
+	MemoryContext oldcxt = MemoryContextSwitchTo(state->context);
+
 	/*
 	 * Copy the tuple.	(Must do this even in WRITEFILE case.)
 	 */
 	tuple = COPYTUP(state, tuple);
 
 	tuplestore_puttuple_common(state, (void *) tuple);
+
+	MemoryContextSwitchTo(oldcxt);
 }
 
 /*
@@ -568,10 +580,13 @@ tuplestore_putvalues(Tuplestorestate *state, TupleDesc tdesc,
 					 Datum *values, bool *isnull)
 {
 	MinimalTuple tuple;
+	MemoryContext oldcxt = MemoryContextSwitchTo(state->context);
 
 	tuple = heap_form_minimal_tuple(tdesc, values, isnull);
 
 	tuplestore_puttuple_common(state, (void *) tuple);
+
+	MemoryContextSwitchTo(oldcxt);	
 }
 
 static void
@@ -579,6 +594,7 @@ tuplestore_puttuple_common(Tuplestorestate *state, void *tuple)
 {
 	TSReadPointer *readptr;
 	int			i;
+	ResourceOwner oldowner;
 
 	switch (state->status)
 	{
@@ -635,8 +651,15 @@ tuplestore_puttuple_common(Tuplestorestate *state, void *tuple)
 			 * the temp file(s) are created in suitable temp tablespaces.
 			 */
 			PrepareTempTablespaces();
+
+			/* associate the file with the store's resource owner */
+			oldowner = CurrentResourceOwner;
+			CurrentResourceOwner = state->resowner;
+
 			state->myfile = BufFileCreateTemp(state->interXact);
 
+			CurrentResourceOwner = oldowner;
+
 			/*
 			 * Freeze the decision about whether trailing length words will be
 			 * used.  We can't change this choice once data is on tape, even