Restore the portal-level snapshot after procedure COMMIT/ROLLBACK.

src/backend/commands/functioncmds.c

@@ -61,6 +61,7 @@
 #include "parser/parse_func.h"
 #include "parser/parse_type.h"
 #include "pgstat.h"
+#include "tcop/pquery.h"
 #include "utils/acl.h"
 #include "utils/builtins.h"
 #include "utils/fmgroids.h"
@@ -2342,6 +2343,20 @@ ExecuteCallStmt(CallStmt *stmt, ParamListInfo params, bool atomic, DestReceiver
 		if (fcinfo.isnull)
 			elog(ERROR, "procedure returned null record");
 
+		/*
+		 * Ensure there's an active snapshot whilst we execute whatever's
+		 * involved here.  Note that this is *not* sufficient to make the
+		 * world safe for TOAST pointers to be included in the returned data:
+		 * the referenced data could have gone away while we didn't hold a
+		 * snapshot.  Hence, it's incumbent on PLs that can do COMMIT/ROLLBACK
+		 * to not return TOAST pointers, unless those pointers were fetched
+		 * after the last COMMIT/ROLLBACK in the procedure.
+		 *
+		 * XXX that is a really nasty, hard-to-test requirement.  Is there a
+		 * way to remove it?
+		 */
+		EnsurePortalSnapshotExists();
+
 		td = DatumGetHeapTupleHeader(retval);
 		tupType = HeapTupleHeaderGetTypeId(td);
 		tupTypmod = HeapTupleHeaderGetTypMod(td);

src/backend/executor/spi.c

@@ -260,12 +260,8 @@ SPI_commit(void)
 	/* Start the actual commit */
 	_SPI_current->internal_xact = true;
 
-	/*
-	 * Before committing, pop all active snapshots to avoid error about
-	 * "snapshot %p still active".
-	 */
-	while (ActiveSnapshotSet())
-		PopActiveSnapshot();
+	/* Release snapshots associated with portals */
+	ForgetPortalSnapshots();
 
 	CommitTransactionCommand();
 	MemoryContextSwitchTo(oldcontext);
@@ -300,6 +296,9 @@ SPI_rollback(void)
 	/* Start the actual rollback */
 	_SPI_current->internal_xact = true;
 
+	/* Release snapshots associated with portals */
+	ForgetPortalSnapshots();
+
 	AbortCurrentTransaction();
 	MemoryContextSwitchTo(oldcontext);
 
@@ -2102,6 +2101,7 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 	Oid			my_lastoid = InvalidOid;
 	SPITupleTable *my_tuptable = NULL;
 	int			res = 0;
+	bool		allow_nonatomic = plan->no_snapshots;	/* legacy API name */
 	bool		pushed_active_snap = false;
 	ErrorContextCallback spierrcontext;
 	CachedPlan *cplan = NULL;
@@ -2134,11 +2134,12 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 	 * In the first two cases, we can just push the snap onto the stack once
 	 * for the whole plan list.
 	 *
-	 * But if the plan has no_snapshots set to true, then don't manage
-	 * snapshots at all.  The caller should then take care of that.
+	 * Note that snapshot != InvalidSnapshot implies an atomic execution
+	 * context.
 	 */
-	if (snapshot != InvalidSnapshot && !plan->no_snapshots)
+	if (snapshot != InvalidSnapshot)
 	{
+		Assert(!allow_nonatomic);
 		if (read_only)
 		{
 			PushActiveSnapshot(snapshot);
@@ -2225,15 +2226,39 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 		stmt_list = cplan->stmt_list;
 
 		/*
-		 * In the default non-read-only case, get a new snapshot, replacing
-		 * any that we pushed in a previous cycle.
+		 * If we weren't given a specific snapshot to use, and the statement
+		 * list requires a snapshot, set that up.
 		 */
-		if (snapshot == InvalidSnapshot && !read_only && !plan->no_snapshots)
+		if (snapshot == InvalidSnapshot &&
+			(list_length(stmt_list) > 1 ||
+			 (list_length(stmt_list) == 1 &&
+			  PlannedStmtRequiresSnapshot(linitial_node(PlannedStmt,
+														stmt_list)))))
 		{
-			if (pushed_active_snap)
-				PopActiveSnapshot();
-			PushActiveSnapshot(GetTransactionSnapshot());
-			pushed_active_snap = true;
+			/*
+			 * First, ensure there's a Portal-level snapshot.  This back-fills
+			 * the snapshot stack in case the previous operation was a COMMIT
+			 * or ROLLBACK inside a procedure or DO block.  (We can't put back
+			 * the Portal snapshot any sooner, or we'd break cases like doing
+			 * SET or LOCK just after COMMIT.)  It's enough to check once per
+			 * statement list, since COMMIT/ROLLBACK/CALL/DO can't appear
+			 * within a multi-statement list.
+			 */
+			EnsurePortalSnapshotExists();
+
+			/*
+			 * In the default non-read-only case, get a new per-statement-list
+			 * snapshot, replacing any that we pushed in a previous cycle.
+			 * Skip it when doing non-atomic execution, though (we rely
+			 * entirely on the Portal snapshot in that case).
+			 */
+			if (!read_only && !allow_nonatomic)
+			{
+				if (pushed_active_snap)
+					PopActiveSnapshot();
+				PushActiveSnapshot(GetTransactionSnapshot());
+				pushed_active_snap = true;
+			}
 		}
 
 		foreach(lc2, stmt_list)
@@ -2246,6 +2271,7 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 			_SPI_current->lastoid = InvalidOid;
 			_SPI_current->tuptable = NULL;
 
+			/* Check for unsupported cases. */
 			if (stmt->utilityStmt)
 			{
 				if (IsA(stmt->utilityStmt, CopyStmt))
@@ -2277,9 +2303,10 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 
 			/*
 			 * If not read-only mode, advance the command counter before each
-			 * command and update the snapshot.
+			 * command and update the snapshot.  (But skip it if the snapshot
+			 * isn't under our control.)
 			 */
-			if (!read_only && !plan->no_snapshots)
+			if (!read_only && pushed_active_snap)
 			{
 				CommandCounterIncrement();
 				UpdateActiveSnapshotCommandId();
@@ -2313,13 +2340,11 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 				ProcessUtilityContext context;
 
 				/*
-				 * If the SPI context is atomic, or we are asked to manage
-				 * snapshots, then we are in an atomic execution context.
-				 * Conversely, to propagate a nonatomic execution context, the
-				 * caller must be in a nonatomic SPI context and manage
-				 * snapshots itself.
+				 * If the SPI context is atomic, or we were not told to allow
+				 * nonatomic operations, tell ProcessUtility this is an atomic
+				 * execution context.
 				 */
-				if (_SPI_current->atomic || !plan->no_snapshots)
+				if (_SPI_current->atomic || !allow_nonatomic)
 					context = PROCESS_UTILITY_QUERY;
 				else
 					context = PROCESS_UTILITY_QUERY_NONATOMIC;

src/backend/replication/logical/worker.c

@@ -194,6 +194,13 @@ create_estate_for_relation(LogicalRepRelMapEntry *rel)
 	ResultRelInfo *resultRelInfo;
 	RangeTblEntry *rte;
 
+	/*
+	 * Input functions may need an active snapshot, as may AFTER triggers
+	 * invoked during finish_estate.  For safety, ensure an active snapshot
+	 * exists throughout all our usage of the executor.
+	 */
+	PushActiveSnapshot(GetTransactionSnapshot());
+
 	estate = CreateExecutorState();
 
 	rte = makeNode(RangeTblEntry);
@@ -221,6 +228,22 @@ create_estate_for_relation(LogicalRepRelMapEntry *rel)
 	return estate;
 }
 
+/*
+ * Finish any operations related to the executor state created by
+ * create_estate_for_relation().
+ */
+static void
+finish_estate(EState *estate)
+{
+	/* Handle any queued AFTER triggers. */
+	AfterTriggerEndQuery(estate);
+
+	/* Cleanup. */
+	ExecResetTupleTable(estate->es_tupleTable, false);
+	FreeExecutorState(estate);
+	PopActiveSnapshot();
+}
+
 /*
  * Executes default values for columns for which we can't map to remote
  * relation columns.
@@ -627,9 +650,6 @@ apply_handle_insert(StringInfo s)
 	remoteslot = ExecInitExtraTupleSlot(estate,
 										RelationGetDescr(rel->localrel));
 
-	/* Input functions may need an active snapshot, so get one */
-	PushActiveSnapshot(GetTransactionSnapshot());
-
 	/* Process and store remote tuple in the slot */
 	oldctx = MemoryContextSwitchTo(GetPerTupleMemoryContext(estate));
 	slot_store_cstrings(remoteslot, rel, newtup.values);
@@ -643,13 +663,8 @@ apply_handle_insert(StringInfo s)
 
 	/* Cleanup. */
 	ExecCloseIndices(estate->es_result_relation_info);
-	PopActiveSnapshot();
 
-	/* Handle queued AFTER triggers. */
-	AfterTriggerEndQuery(estate);
-
-	ExecResetTupleTable(estate->es_tupleTable, false);
-	FreeExecutorState(estate);
+	finish_estate(estate);
 
 	logicalrep_rel_close(rel, NoLock);
 
@@ -760,7 +775,6 @@ apply_handle_update(StringInfo s)
 		}
 	}
 
-	PushActiveSnapshot(GetTransactionSnapshot());
 	ExecOpenIndices(estate->es_result_relation_info, false);
 
 	/* Build the search tuple. */
@@ -819,15 +833,10 @@ apply_handle_update(StringInfo s)
 	}
 
 	/* Cleanup. */
+	EvalPlanQualEnd(&epqstate);
 	ExecCloseIndices(estate->es_result_relation_info);
-	PopActiveSnapshot();
 
-	/* Handle queued AFTER triggers. */
-	AfterTriggerEndQuery(estate);
-
-	EvalPlanQualEnd(&epqstate);
-	ExecResetTupleTable(estate->es_tupleTable, false);
-	FreeExecutorState(estate);
+	finish_estate(estate);
 
 	logicalrep_rel_close(rel, NoLock);
 
@@ -878,7 +887,6 @@ apply_handle_delete(StringInfo s)
 									   RelationGetDescr(rel->localrel));
 	EvalPlanQualInit(&epqstate, estate, NULL, NIL, -1);
 
-	PushActiveSnapshot(GetTransactionSnapshot());
 	ExecOpenIndices(estate->es_result_relation_info, false);
 
 	/* Find the tuple using the replica identity index. */
@@ -919,15 +927,10 @@ apply_handle_delete(StringInfo s)
 	}
 
 	/* Cleanup. */
+	EvalPlanQualEnd(&epqstate);
 	ExecCloseIndices(estate->es_result_relation_info);
-	PopActiveSnapshot();
 
-	/* Handle queued AFTER triggers. */
-	AfterTriggerEndQuery(estate);
-
-	EvalPlanQualEnd(&epqstate);
-	ExecResetTupleTable(estate->es_tupleTable, false);
-	FreeExecutorState(estate);
+	finish_estate(estate);
 
 	logicalrep_rel_close(rel, NoLock);