Re: [BUG] failed assertion in EnsurePortalSnapshotExists()

"Drouvot, Bertrand" <bdrouvot@amazon.com> writes:
> [ v2-0003-EnsurePortalSnapshotExists-failed-assertion.patch ]

Looking through this, I think you were overenthusiastic about applying
PushActiveSnapshotWithLevel.  We don't really need to use it except in
the places where we're setting portalSnapshot, because other than those
cases we don't have a risk of portalSnapshot becoming a dangling pointer.
Also, I'm quite nervous about applying PushActiveSnapshotWithLevel to
snapshots that aren't created by the portal machinery itself, because
we don't know very much about where passed-in snapshots came from or
what the caller thinks their lifespan is.

The attached revision therefore backs off to only using the new code
in the two places where we really need it.  I made a number of
more-cosmetic revisions too.  Notably, I think it's useful to frame
the testing shortcoming as "we were not testing COMMIT/ROLLBACK
inside a plpgsql exception block".  So I moved the test code to the
plpgsql tests and made it check ROLLBACK too.

            regards, tom lane

PS: Memo to self: in the back branches, the new field has to be
added at the end of struct Portal.

diff --git a/src/backend/access/transam/xact.c b/src/backend/access/transam/xact.c
index 6597ec45a9..4cc38f0d85 100644
--- a/src/backend/access/transam/xact.c
+++ b/src/backend/access/transam/xact.c
@@ -4863,6 +4863,7 @@ CommitSubTransaction(void)
     AfterTriggerEndSubXact(true);
     AtSubCommit_Portals(s->subTransactionId,
                         s->parent->subTransactionId,
+                        s->parent->nestingLevel,
                         s->parent->curTransactionOwner);
     AtEOSubXact_LargeObject(true, s->subTransactionId,
                             s->parent->subTransactionId);
diff --git a/src/backend/tcop/pquery.c b/src/backend/tcop/pquery.c
index b5797042ab..960f3fadce 100644
--- a/src/backend/tcop/pquery.c
+++ b/src/backend/tcop/pquery.c
@@ -480,7 +480,9 @@ PortalStart(Portal portal, ParamListInfo params,
                  * We could remember the snapshot in portal->portalSnapshot,
                  * but presently there seems no need to, as this code path
                  * cannot be used for non-atomic execution.  Hence there can't
-                 * be any commit/abort that might destroy the snapshot.
+                 * be any commit/abort that might destroy the snapshot.  Since
+                 * we don't do that, there's also no need to force a
+                 * non-default nesting level for the snapshot.
                  */

                 /*
@@ -1136,9 +1138,15 @@ PortalRunUtility(Portal portal, PlannedStmt *pstmt,
             snapshot = RegisterSnapshot(snapshot);
             portal->holdSnapshot = snapshot;
         }
-        /* In any case, make the snapshot active and remember it in portal */
-        PushActiveSnapshot(snapshot);
-        /* PushActiveSnapshot might have copied the snapshot */
+
+        /*
+         * In any case, make the snapshot active and remember it in portal.
+         * Because the portal now references the snapshot, we must tell
+         * snapmgr.c that the snapshot belongs to the portal's transaction
+         * level, else we risk portalSnapshot becoming a dangling pointer.
+         */
+        PushActiveSnapshotWithLevel(snapshot, portal->createLevel);
+        /* PushActiveSnapshotWithLevel might have copied the snapshot */
         portal->portalSnapshot = GetActiveSnapshot();
     }
     else
@@ -1784,8 +1792,13 @@ EnsurePortalSnapshotExists(void)
         elog(ERROR, "cannot execute SQL without an outer snapshot or portal");
     Assert(portal->portalSnapshot == NULL);

-    /* Create a new snapshot and make it active */
-    PushActiveSnapshot(GetTransactionSnapshot());
-    /* PushActiveSnapshot might have copied the snapshot */
+    /*
+     * Create a new snapshot, make it active, and remember it in portal.
+     * Because the portal now references the snapshot, we must tell snapmgr.c
+     * that the snapshot belongs to the portal's transaction level, else we
+     * risk portalSnapshot becoming a dangling pointer.
+     */
+    PushActiveSnapshotWithLevel(GetTransactionSnapshot(), portal->createLevel);
+    /* PushActiveSnapshotWithLevel might have copied the snapshot */
     portal->portalSnapshot = GetActiveSnapshot();
 }
diff --git a/src/backend/utils/mmgr/portalmem.c b/src/backend/utils/mmgr/portalmem.c
index 5c30e141f5..58674d611d 100644
--- a/src/backend/utils/mmgr/portalmem.c
+++ b/src/backend/utils/mmgr/portalmem.c
@@ -210,6 +210,7 @@ CreatePortal(const char *name, bool allowDup, bool dupSilent)
     portal->cleanup = PortalCleanup;
     portal->createSubid = GetCurrentSubTransactionId();
     portal->activeSubid = portal->createSubid;
+    portal->createLevel = GetCurrentTransactionNestLevel();
     portal->strategy = PORTAL_MULTI_QUERY;
     portal->cursorOptions = CURSOR_OPT_NO_SCROLL;
     portal->atStart = true;
@@ -657,6 +658,7 @@ HoldPortal(Portal portal)
      */
     portal->createSubid = InvalidSubTransactionId;
     portal->activeSubid = InvalidSubTransactionId;
+    portal->createLevel = 0;
 }

 /*
@@ -940,6 +942,7 @@ PortalErrorCleanup(void)
 void
 AtSubCommit_Portals(SubTransactionId mySubid,
                     SubTransactionId parentSubid,
+                    int parentLevel,
                     ResourceOwner parentXactOwner)
 {
     HASH_SEQ_STATUS status;
@@ -954,6 +957,7 @@ AtSubCommit_Portals(SubTransactionId mySubid,
         if (portal->createSubid == mySubid)
         {
             portal->createSubid = parentSubid;
+            portal->createLevel = parentLevel;
             if (portal->resowner)
                 ResourceOwnerNewParent(portal->resowner, parentXactOwner);
         }
diff --git a/src/backend/utils/time/snapmgr.c b/src/backend/utils/time/snapmgr.c
index 2968c7f7b7..dca1bc8afc 100644
--- a/src/backend/utils/time/snapmgr.c
+++ b/src/backend/utils/time/snapmgr.c
@@ -678,10 +678,25 @@ FreeSnapshot(Snapshot snapshot)
  */
 void
 PushActiveSnapshot(Snapshot snap)
+{
+    PushActiveSnapshotWithLevel(snap, GetCurrentTransactionNestLevel());
+}
+
+/*
+ * PushActiveSnapshotWithLevel
+ *        Set the given snapshot as the current active snapshot
+ *
+ * Same as PushActiveSnapshot except that caller can specify the
+ * transaction nesting level that "owns" the snapshot.  This level
+ * must not be deeper than the current top of the snapshot stack.
+ */
+void
+PushActiveSnapshotWithLevel(Snapshot snap, int snap_level)
 {
     ActiveSnapshotElt *newactive;

     Assert(snap != InvalidSnapshot);
+    Assert(ActiveSnapshot == NULL || snap_level >= ActiveSnapshot->as_level);

     newactive = MemoryContextAlloc(TopTransactionContext, sizeof(ActiveSnapshotElt));

@@ -695,7 +710,7 @@ PushActiveSnapshot(Snapshot snap)
         newactive->as_snap = snap;

     newactive->as_next = ActiveSnapshot;
-    newactive->as_level = GetCurrentTransactionNestLevel();
+    newactive->as_level = snap_level;

     newactive->as_snap->active_count++;

diff --git a/src/include/utils/portal.h b/src/include/utils/portal.h
index 2e5bbdd0ec..37b1817ae9 100644
--- a/src/include/utils/portal.h
+++ b/src/include/utils/portal.h
@@ -130,6 +130,7 @@ typedef struct PortalData
      */
     SubTransactionId createSubid;    /* the creating subxact */
     SubTransactionId activeSubid;    /* the last subxact with activity */
+    int            createLevel;    /* creating subxact's nesting level */

     /* The query or queries the portal will execute */
     const char *sourceText;        /* text of query (as of 8.4, never NULL) */
@@ -219,6 +220,7 @@ extern void AtCleanup_Portals(void);
 extern void PortalErrorCleanup(void);
 extern void AtSubCommit_Portals(SubTransactionId mySubid,
                                 SubTransactionId parentSubid,
+                                int parentLevel,
                                 ResourceOwner parentXactOwner);
 extern void AtSubAbort_Portals(SubTransactionId mySubid,
                                SubTransactionId parentSubid,
diff --git a/src/include/utils/snapmgr.h b/src/include/utils/snapmgr.h
index 44539fe15a..c6a176cc95 100644
--- a/src/include/utils/snapmgr.h
+++ b/src/include/utils/snapmgr.h
@@ -114,6 +114,7 @@ extern void InvalidateCatalogSnapshot(void);
 extern void InvalidateCatalogSnapshotConditionally(void);

 extern void PushActiveSnapshot(Snapshot snapshot);
+extern void PushActiveSnapshotWithLevel(Snapshot snapshot, int snap_level);
 extern void PushCopiedSnapshot(Snapshot snapshot);
 extern void UpdateActiveSnapshotCommandId(void);
 extern void PopActiveSnapshot(void);
diff --git a/src/pl/plpgsql/src/expected/plpgsql_transaction.out b/src/pl/plpgsql/src/expected/plpgsql_transaction.out
index f79f847321..254e5b7a70 100644
--- a/src/pl/plpgsql/src/expected/plpgsql_transaction.out
+++ b/src/pl/plpgsql/src/expected/plpgsql_transaction.out
@@ -430,6 +430,34 @@ SELECT * FROM test1;
 ---+---
 (0 rows)

+-- test commit/rollback inside exception handler, too
+TRUNCATE test1;
+DO LANGUAGE plpgsql $$
+BEGIN
+    FOR i IN 1..10 LOOP
+      BEGIN
+        INSERT INTO test1 VALUES (i, 'good');
+        INSERT INTO test1 VALUES (i/0, 'bad');
+      EXCEPTION
+        WHEN division_by_zero THEN
+            INSERT INTO test1 VALUES (i, 'exception');
+            IF (i % 3) > 0 THEN COMMIT; ELSE ROLLBACK; END IF;
+      END;
+    END LOOP;
+END;
+$$;
+SELECT * FROM test1;
+ a  |     b
+----+-----------
+  1 | exception
+  2 | exception
+  4 | exception
+  5 | exception
+  7 | exception
+  8 | exception
+ 10 | exception
+(7 rows)
+
 -- detoast result of simple expression after commit
 CREATE TEMP TABLE test4(f1 text);
 ALTER TABLE test4 ALTER COLUMN f1 SET STORAGE EXTERNAL; -- disable compression
diff --git a/src/pl/plpgsql/src/sql/plpgsql_transaction.sql b/src/pl/plpgsql/src/sql/plpgsql_transaction.sql
index 888ddccace..8d76d00daa 100644
--- a/src/pl/plpgsql/src/sql/plpgsql_transaction.sql
+++ b/src/pl/plpgsql/src/sql/plpgsql_transaction.sql
@@ -354,6 +354,27 @@ $$;
 SELECT * FROM test1;


+-- test commit/rollback inside exception handler, too
+TRUNCATE test1;
+
+DO LANGUAGE plpgsql $$
+BEGIN
+    FOR i IN 1..10 LOOP
+      BEGIN
+        INSERT INTO test1 VALUES (i, 'good');
+        INSERT INTO test1 VALUES (i/0, 'bad');
+      EXCEPTION
+        WHEN division_by_zero THEN
+            INSERT INTO test1 VALUES (i, 'exception');
+            IF (i % 3) > 0 THEN COMMIT; ELSE ROLLBACK; END IF;
+      END;
+    END LOOP;
+END;
+$$;
+
+SELECT * FROM test1;
+
+
 -- detoast result of simple expression after commit
 CREATE TEMP TABLE test4(f1 text);
 ALTER TABLE test4 ALTER COLUMN f1 SET STORAGE EXTERNAL; -- disable compression