Preserve CurrentMemoryContext across Start/CommitTransactionCommand.

tglsfdc · tglsfdc · commit 1afe31f03cd2 · 2024-07-01T11:55:19.000-04:00
Up to now, committing a transaction has caused CurrentMemoryContext to get set to TopMemoryContext. Most callers did not pay any particular heed to this, which is problematic because TopMemoryContext is a long-lived context that never gets reset. If the caller assumes it can leak memory because it's running in a limited-lifespan context, that behavior translates into a session-lifespan memory leak. The first-reported instance of this involved ProcessIncomingNotify, which is called from the main processing loop that normally runs in MessageContext. That outer-loop code assumes that whatever it allocates will be cleaned up when we're done processing the current client message --- but if we service a notify interrupt, then whatever gets allocated before the next switch to MessageContext will be permanently leaked in TopMemoryContext. sinval catchup interrupts have a similar problem, and I strongly suspect that some places in logical replication do too. To fix this in a generic way, let's redefine the behavior as "CommitTransactionCommand restores the memory context that was current at entry to StartTransactionCommand". This clearly fixes the issue for the notify and sinval cases, and it seems to match the mental model that's in use in the logical replication code, to the extent that anybody thought about it there at all. For consistency, likewise make subtransaction exit restore the context that was current at subtransaction start (rather than always selecting the CurTransactionContext of the parent transaction level). This case has less risk of resulting in a permanent leak than the outer-level behavior has, but it would not meet the principle of least surprise for some CommitTransactionCommand calls to restore the previous context while others don't. While we're here, also change xact.c so that we reset TopTransactionContext at transaction exit and then re-use it in later transactions, rather than dropping and recreating it in each cycle. This probably doesn't save a lot given the context recycling mechanism in aset.c, but it should save a little bit. Per suggestion from David Rowley. (Parenthetically, the text in src/backend/utils/mmgr/README implies that this is how I'd planned to implement it as far back as commit 1aebc36 --- but the code actually added in that commit just drops and recreates it each time.) This commit also cleans up a few places outside xact.c that were needlessly making TopMemoryContext current, and thus risking more leaks of the same kind. I don't think any of them represent significant leak risks today, but let's deal with them while the issue is top-of-mind. Per bug #18512 from wizardbrony. Commit to HEAD only, as this change seems to have some risk of breaking things for some callers. We'll apply a narrower fix for the known-broken cases in the back branches. Discussion: https://postgr.es/m/3478884.1718656625@sss.pgh.pa.us
diff --git 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 struct.)
+	 */
 	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 struct.)
+	 */
 	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
@@ -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
@@ -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
@@ -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
@@ -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"),

Original file line number	Diff line number	Diff line change
`@@ -1148,7 +1148,8 @@ MemoryContextAllocationFailure(MemoryContext context, Size size, int flags)`
`1148`	`1148`	`{`
`1149`	`1149`	`if ((flags & MCXT_ALLOC_NO_OOM) == 0)`
`1150`	`1150`	`{`
`1151`		`- MemoryContextStats(TopMemoryContext);`
	`1151`	`+ if (TopMemoryContext)`
	`1152`	`+ MemoryContextStats(TopMemoryContext);`
`1152`	`1153`	`ereport(ERROR,`
`1153`	`1154`	`(errcode(ERRCODE_OUT_OF_MEMORY),`
`1154`	`1155`	`errmsg("out of memory"),`