Re: Xact end leaves CurrentMemoryContext = TopMemoryContext

David Rowley <dgrowleyml@gmail.com> writes:
> On Tue, 18 Jun 2024 at 16:53, Tom Lane <tgl@sss.pgh.pa.us> wrote:
>> I'll poke at this tomorrow, unless you're hot to try it right now.

> Please go ahead. I was just in suggestion mode here.

So I tried that, and while it kind of worked, certain parts of the
system (notably logical replication) had acute indigestion.  Turns
out there is a fair amount of code that does

    StartTransactionCommand();
    ... some random thing or other ...
    CommitTransactionCommand();

and does not stop to think at all about whether that has any effect on
its memory context.  Andres' idea would break every single such place,
and this idea isn't much better because while it does provide a
current memory context after CommitTransactionCommand, that context is
effectively short-lived: the next Start/CommitTransactionCommand will
trash it.  That broke a lot more places than I'd hoped, mostly in
obscure ways.

After awhile I had an epiphany: what we should do is make
CommitTransactionCommand restore the memory context that was active
before StartTransactionCommand.  That's what we want in every place
that was cognizant of this issue, and it seems to be the case in every
place that wasn't doing anything explicit about it, either.

The 0001 patch attached does that, and seems to work nicely.
I made it implement the idea of recycling TopTransactionContext,
too.  (Interestingly, src/backend/utils/mmgr/README *already*
claims we manage TopTransactionContext this way.  Did we do that
and then change it back in the mists of time?)  The core parts
of the patch are all in xact.c --- the other diffs are just random
cleanup that I found while surveying use of TopMemoryContext and
CommitTransactionCommand.

Also, 0002 is what I propose for the back branches.  It just adds
memory context save/restore in notify interrupt processing to solve
the original bug report, as well as in sinval interrupt processing
which I discovered has the same disease.  We don't need this in
HEAD if we apply 0001.

At this point I'd be inclined to wait for the branch to be made,
and then apply 0001 in HEAD/v18 only and 0002 in v17 and before.
While 0001 seems fairly straightforward, it's still a nontrivial
change and I'm hesitant to shove it in at this late stage of the
v17 cycle.

            regards, tom lane

diff --git a/src/backend/access/transam/xact.c b/src/backend/access/transam/xact.c
index 9bda1aa6bc..fce12fea6e 100644
--- a/src/backend/access/transam/xact.c
+++ b/src/backend/access/transam/xact.c
@@ -200,6 +200,7 @@ typedef struct TransactionStateData
     int            gucNestLevel;    /* GUC context nesting depth */
     MemoryContext curTransactionContext;    /* my xact-lifetime context */
     ResourceOwner curTransactionOwner;    /* my query resources */
+    MemoryContext priorContext; /* CurrentMemoryContext before xact started */
     TransactionId *childXids;    /* subcommitted child XIDs, in XID order */
     int            nChildXids;        /* # of subcommitted child XIDs */
     int            maxChildXids;    /* allocated size of childXids[] */
@@ -1174,6 +1175,11 @@ AtStart_Memory(void)
 {
     TransactionState s = CurrentTransactionState;

+    /*
+     * Remember the memory context that was active prior to transaction start.
+     */
+    s->priorContext = CurrentMemoryContext;
+
     /*
      * If this is the first time through, create a private context for
      * AbortTransaction to work in.  By reserving some space now, we can
@@ -1190,17 +1196,15 @@ AtStart_Memory(void)
                                   32 * 1024);

     /*
-     * We shouldn't have a transaction context already.
-     */
-    Assert(TopTransactionContext == NULL);
-
-    /*
-     * Create a toplevel context for the transaction.
+     * Likewise, if this is the first time through, create a top-level context
+     * for transaction-local data.  This context will be reset at transaction
+     * end, and then re-used in later transactions.
      */
-    TopTransactionContext =
-        AllocSetContextCreate(TopMemoryContext,
-                              "TopTransactionContext",
-                              ALLOCSET_DEFAULT_SIZES);
+    if (TopTransactionContext == NULL)
+        TopTransactionContext =
+            AllocSetContextCreate(TopMemoryContext,
+                                  "TopTransactionContext",
+                                  ALLOCSET_DEFAULT_SIZES);

     /*
      * In a top-level transaction, CurTransactionContext is the same as
@@ -1251,6 +1255,11 @@ AtSubStart_Memory(void)

     Assert(CurTransactionContext != NULL);

+    /*
+     * Remember the context that was active prior to subtransaction start.
+     */
+    s->priorContext = CurrentMemoryContext;
+
     /*
      * Create a CurTransactionContext, which will be used to hold data that
      * survives subtransaction commit but disappears on subtransaction abort.
@@ -1576,20 +1585,30 @@ AtCCI_LocalCache(void)
 static void
 AtCommit_Memory(void)
 {
+    TransactionState s = CurrentTransactionState;
+
     /*
-     * Now that we're "out" of a transaction, have the system allocate things
-     * in the top memory context instead of per-transaction contexts.
+     * Return to the memory context that was current before we started the
+     * transaction.  (In principle, this could not be any of the contexts we
+     * are about to delete.  If it somehow is, assertions in mcxt.c will
+     * complain.)
      */
-    MemoryContextSwitchTo(TopMemoryContext);
+    MemoryContextSwitchTo(s->priorContext);

     /*
-     * Release all transaction-local memory.
+     * Release all transaction-local memory.  TopTransactionContext survives
+     * but becomes empty; any sub-contexts go away.
      */
     Assert(TopTransactionContext != NULL);
-    MemoryContextDelete(TopTransactionContext);
-    TopTransactionContext = NULL;
+    MemoryContextReset(TopTransactionContext);
+
+    /*
+     * Clear these pointers as a pro-forma matter.  (Notionally, while
+     * TopTransactionContext still exists, it's currently not associated with
+     * this TransactionState block.)
+     */
     CurTransactionContext = NULL;
-    CurrentTransactionState->curTransactionContext = NULL;
+    s->curTransactionContext = NULL;
 }

 /* ----------------------------------------------------------------
@@ -1942,13 +1961,18 @@ AtSubAbort_childXids(void)
 static void
 AtCleanup_Memory(void)
 {
-    Assert(CurrentTransactionState->parent == NULL);
+    TransactionState s = CurrentTransactionState;
+
+    /* Should be at top level */
+    Assert(s->parent == NULL);

     /*
-     * Now that we're "out" of a transaction, have the system allocate things
-     * in the top memory context instead of per-transaction contexts.
+     * Return to the memory context that was current before we started the
+     * transaction.  (In principle, this could not be any of the contexts we
+     * are about to delete.  If it somehow is, assertions in mcxt.c will
+     * complain.)
      */
-    MemoryContextSwitchTo(TopMemoryContext);
+    MemoryContextSwitchTo(s->priorContext);

     /*
      * Clear the special abort context for next time.
@@ -1957,13 +1981,20 @@ AtCleanup_Memory(void)
         MemoryContextReset(TransactionAbortContext);

     /*
-     * Release all transaction-local memory.
+     * Release all transaction-local memory, the same as in AtCommit_Memory,
+     * except we must cope with the possibility that we didn't get as far as
+     * creating TopTransactionContext.
      */
     if (TopTransactionContext != NULL)
-        MemoryContextDelete(TopTransactionContext);
-    TopTransactionContext = NULL;
+        MemoryContextReset(TopTransactionContext);
+
+    /*
+     * Clear these pointers as a pro-forma matter.  (Notionally, while
+     * TopTransactionContext still exists, it's currently not associated with
+     * this TransactionState block.)
+     */
     CurTransactionContext = NULL;
-    CurrentTransactionState->curTransactionContext = NULL;
+    s->curTransactionContext = NULL;
 }


@@ -1982,8 +2013,15 @@ AtSubCleanup_Memory(void)

     Assert(s->parent != NULL);

-    /* Make sure we're not in an about-to-be-deleted context */
-    MemoryContextSwitchTo(s->parent->curTransactionContext);
+    /*
+     * Return to the memory context that was current before we started the
+     * subtransaction.  (In principle, this could not be any of the contexts
+     * we are about to delete.  If it somehow is, assertions in mcxt.c will
+     * complain.)
+     */
+    MemoryContextSwitchTo(s->priorContext);
+
+    /* Update CurTransactionContext (might not be same as priorContext) */
     CurTransactionContext = s->parent->curTransactionContext;

     /*
@@ -4904,8 +4942,13 @@ AbortOutOfAnyTransaction(void)
     /* Should be out of all subxacts now */
     Assert(s->parent == NULL);

-    /* If we didn't actually have anything to do, revert to TopMemoryContext */
-    AtCleanup_Memory();
+    /*
+     * Revert to TopMemoryContext, to ensure we exit in a well-defined state
+     * whether there were any transactions to close or not.  (Callers that
+     * don't intend to exit soon should switch to some other context to avoid
+     * long-term memory leaks.)
+     */
+    MemoryContextSwitchTo(TopMemoryContext);
 }

 /*
diff --git a/src/backend/replication/walsender.c b/src/backend/replication/walsender.c
index c623b07cf0..e78a5d51cf 100644
--- a/src/backend/replication/walsender.c
+++ b/src/backend/replication/walsender.c
@@ -440,12 +440,10 @@ IdentifySystem(void)

         /* syscache access needs a transaction env. */
         StartTransactionCommand();
-        /* make dbname live outside TX context */
-        MemoryContextSwitchTo(cur);
         dbname = get_database_name(MyDatabaseId);
+        /* copy dbname out of TX context */
+        dbname = MemoryContextStrdup(cur, dbname);
         CommitTransactionCommand();
-        /* CommitTransactionCommand switches to TopMemoryContext */
-        MemoryContextSwitchTo(cur);
     }

     dest = CreateDestReceiver(DestRemoteSimple);
diff --git a/src/backend/tcop/backend_startup.c b/src/backend/tcop/backend_startup.c
index cfa2755196..a2f94b1050 100644
--- a/src/backend/tcop/backend_startup.c
+++ b/src/backend/tcop/backend_startup.c
@@ -196,9 +196,8 @@ BackendInitialize(ClientSocket *client_sock, CAC_state cac)
      * Save remote_host and remote_port in port structure (after this, they
      * will appear in log_line_prefix data for log messages).
      */
-    oldcontext = MemoryContextSwitchTo(TopMemoryContext);
-    port->remote_host = pstrdup(remote_host);
-    port->remote_port = pstrdup(remote_port);
+    port->remote_host = MemoryContextStrdup(TopMemoryContext, remote_host);
+    port->remote_port = MemoryContextStrdup(TopMemoryContext, remote_port);

     /* And now we can issue the Log_connections message, if wanted */
     if (Log_connections)
@@ -230,9 +229,8 @@ BackendInitialize(ClientSocket *client_sock, CAC_state cac)
         strspn(remote_host, "0123456789.") < strlen(remote_host) &&
         strspn(remote_host, "0123456789ABCDEFabcdef:") < strlen(remote_host))
     {
-        port->remote_hostname = pstrdup(remote_host);
+        port->remote_hostname = MemoryContextStrdup(TopMemoryContext, remote_host);
     }
-    MemoryContextSwitchTo(oldcontext);

     /*
      * Ready to begin client interaction.  We will give up and _exit(1) after
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index 45a3794b8e..adf71c6902 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -4418,7 +4418,7 @@ PostgresMain(const char *dbname, const char *username)
          * Now return to normal top-level context and clear ErrorContext for
          * next time.
          */
-        MemoryContextSwitchTo(TopMemoryContext);
+        MemoryContextSwitchTo(MessageContext);
         FlushErrorState();

         /*
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index b42ececbe7..bde54326c6 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -1148,7 +1148,8 @@ MemoryContextAllocationFailure(MemoryContext context, Size size, int flags)
 {
     if ((flags & MCXT_ALLOC_NO_OOM) == 0)
     {
-        MemoryContextStats(TopMemoryContext);
+        if (TopMemoryContext)
+            MemoryContextStats(TopMemoryContext);
         ereport(ERROR,
                 (errcode(ERRCODE_OUT_OF_MEMORY),
                  errmsg("out of memory"),
diff --git a/src/backend/commands/async.c b/src/backend/commands/async.c
index ab4c72762d..648ad4a846 100644
--- a/src/backend/commands/async.c
+++ b/src/backend/commands/async.c
@@ -2182,6 +2182,8 @@ asyncQueueAdvanceTail(void)
 static void
 ProcessIncomingNotify(bool flush)
 {
+    MemoryContext oldcontext;
+
     /* We *must* reset the flag */
     notifyInterruptPending = false;

@@ -2196,14 +2198,21 @@ ProcessIncomingNotify(bool flush)

     /*
      * We must run asyncQueueReadAllNotifications inside a transaction, else
-     * bad things happen if it gets an error.
+     * bad things happen if it gets an error.  However, we need to preserve
+     * the caller's memory context (typically MessageContext).
      */
+    oldcontext = CurrentMemoryContext;
+
     StartTransactionCommand();

     asyncQueueReadAllNotifications();

     CommitTransactionCommand();

+    /* Caller's context had better not have been transaction-local */
+    Assert(MemoryContextIsValid(oldcontext));
+    MemoryContextSwitchTo(oldcontext);
+
     /*
      * If this isn't an end-of-command case, we must flush the notify messages
      * to ensure frontend gets them promptly.
diff --git a/src/backend/storage/ipc/sinval.c b/src/backend/storage/ipc/sinval.c
index d9b16f84d1..4e69015fe4 100644
--- a/src/backend/storage/ipc/sinval.c
+++ b/src/backend/storage/ipc/sinval.c
@@ -16,6 +16,7 @@

 #include "access/xact.h"
 #include "miscadmin.h"
+#include "nodes/memnodes.h"
 #include "storage/latch.h"
 #include "storage/sinvaladt.h"
 #include "utils/inval.h"
@@ -182,6 +183,7 @@ ProcessCatchupInterrupt(void)
          * can just call AcceptInvalidationMessages() to do this.  If we
          * aren't, we start and immediately end a transaction; the call to
          * AcceptInvalidationMessages() happens down inside transaction start.
+         * Be sure to preserve caller's memory context when we do that.
          *
          * It is awfully tempting to just call AcceptInvalidationMessages()
          * without the rest of the xact start/stop overhead, and I think that
@@ -195,9 +197,14 @@ ProcessCatchupInterrupt(void)
         }
         else
         {
+            MemoryContext oldcontext = CurrentMemoryContext;
+
             elog(DEBUG4, "ProcessCatchupEvent outside transaction");
             StartTransactionCommand();
             CommitTransactionCommand();
+            /* Caller's context had better not have been transaction-local */
+            Assert(MemoryContextIsValid(oldcontext));
+            MemoryContextSwitchTo(oldcontext);
         }
     }
 }