From 693e9707be2a006c70a06052564c549c37acc07b Mon Sep 17 00:00:00 2001
From: Alexander Korotkov <akorotkov@postgresql.org>
Date: Wed, 6 Mar 2024 13:59:55 +0200
Subject: [PATCH v4 3/3] Avoid recursion in MemoryContext functions.

You might run out of stack space with recursion, which is not nice in
functions that might be used e.g. at cleanup after transaction
abort. MemoryContext contains pointer to parent and siblings, so we
can traverse a tree of contexts iteratively, without using
stack. Refactor the functions to do that.

MemoryContextStats() still recurses, but it now has a limit to how
deep it recurses. Once the limit is reached, it prints just a summary
of the rest of the hierarchy, similar to how it summarizes contexts
with lots of children. That seems good anyway, because a context dump
with hundreds of nested contexts isn't very readable.

Report by Egor Chindyaskin and Alexander Lakhin.

Discussion: https://www.postgresql.org/message-id/1672760457.940462079%40f306.i.mail.ru
Author: Heikki Linnakangas
Reviewed-by: Robert Haas, Andres Freund, Alexander Korotkov
---
 src/backend/utils/mmgr/mcxt.c | 246 +++++++++++++++++++++++-----------
 src/include/utils/memutils.h  |   3 +-
 2 files changed, 170 insertions(+), 79 deletions(-)

diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 1a615becae4..04b540a578e 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -145,9 +145,10 @@ MemoryContext CurTransactionContext = NULL;
 /* This is a transient link to the active portal's memory context: */
 MemoryContext PortalContext = NULL;
 
+static void MemoryContextDeleteOnly(MemoryContext context);
 static void MemoryContextCallResetCallbacks(MemoryContext context);
 static void MemoryContextStatsInternal(MemoryContext context, int level,
-									   bool print, int max_children,
+									   int max_level, int max_children,
 									   MemoryContextCounters *totals,
 									   bool print_to_stderr);
 static void MemoryContextStatsPrint(MemoryContext context, void *passthru,
@@ -256,6 +257,41 @@ BogusGetChunkSpace(void *pointer)
 	return 0;					/* keep compiler quiet */
 }
 
+/*
+ * MemoryContextTraverseNext
+ *		Helper function to traverse all descendants of a memory context
+ *		without recursion.
+ *
+ * Recursion could lead to out-of-stack errors with deep context hierarchies,
+ * which would be unpleasant in error cleanup code paths.
+ *
+ * To process 'context' and all its descendants, use a loop like this:
+ *
+ *     <process 'context'>
+ *     for (MemoryContext curr = context->firstchild;
+ *          curr != NULL;
+ *          curr = MemoryContextTraverseNext(curr, context))
+ *     {
+ *         <process 'curr'>
+ *     }
+ *
+ * This visits all the contexts in the depth-first order.
+ */
+static MemoryContext
+MemoryContextTraverseNext(MemoryContext curr, MemoryContext top)
+{
+	if (curr->firstchild != NULL)
+		return curr->firstchild;
+
+	while (curr->nextchild == NULL)
+	{
+		curr = curr->parent;
+		if (curr == top)
+			return NULL;
+	}
+	return curr->nextchild;
+}
+
 
 /*****************************************************************************
  *	  EXPORTED ROUTINES														 *
@@ -375,14 +411,13 @@ MemoryContextResetOnly(MemoryContext context)
 void
 MemoryContextResetChildren(MemoryContext context)
 {
-	MemoryContext child;
-
 	Assert(MemoryContextIsValid(context));
 
-	for (child = context->firstchild; child != NULL; child = child->nextchild)
+	for (MemoryContext curr = context->firstchild;
+		 curr != NULL;
+		 curr = MemoryContextTraverseNext(curr, context))
 	{
-		MemoryContextResetChildren(child);
-		MemoryContextResetOnly(child);
+		MemoryContextResetOnly(curr);
 	}
 }
 
@@ -397,16 +432,53 @@ MemoryContextResetChildren(MemoryContext context)
  */
 void
 MemoryContextDelete(MemoryContext context)
+{
+	MemoryContext curr;
+
+	/*
+	 * Delete subcontexts from the bottom up.
+	 *
+	 * Note: Do not use recursion here.  A "stack depth limit exceeded" error
+	 * would be unpleasant if we're already in the process of cleaning up from
+	 * transaction abort.  We also cannot use MemoryContextTraverseNext() here
+	 * because we modify the tree as we go.
+	 */
+	curr = context;
+	for (;;)
+	{
+		MemoryContext parent;
+
+		/* Descend down until we find a leaf context with no children */
+		while (curr->firstchild != NULL)
+			curr = curr->firstchild;
+
+		/*
+		 * We're now at a leaf with no children. Free it and continue from the
+		 * parent.  Or if this was the original node, we're all done.
+		 */
+		parent = curr->parent;
+		MemoryContextDeleteOnly(curr);
+
+		if (curr == context)
+			break;
+		curr = parent;
+	}
+}
+
+/*
+ * Subroutine of MemoryContextDelete, to delete a context that has no
+ * children.
+ */
+static void
+MemoryContextDeleteOnly(MemoryContext context)
 {
 	Assert(MemoryContextIsValid(context));
 	/* We had better not be deleting TopMemoryContext ... */
 	Assert(context != TopMemoryContext);
 	/* And not CurrentMemoryContext, either */
 	Assert(context != CurrentMemoryContext);
-
-	/* save a function call in common case where there are no children */
-	if (context->firstchild != NULL)
-		MemoryContextDeleteChildren(context);
+	/* All the children should've been deleted already */
+	Assert(context->firstchild == NULL);
 
 	/*
 	 * It's not entirely clear whether 'tis better to do this before or after
@@ -672,12 +744,12 @@ MemoryContextMemAllocated(MemoryContext context, bool recurse)
 
 	if (recurse)
 	{
-		MemoryContext child;
-
-		for (child = context->firstchild;
-			 child != NULL;
-			 child = child->nextchild)
-			total += MemoryContextMemAllocated(child, true);
+		for (MemoryContext curr = context->firstchild;
+			 curr != NULL;
+			 curr = MemoryContextTraverseNext(curr, context))
+		{
+			total += curr->mem_allocated;
+		}
 	}
 
 	return total;
@@ -691,9 +763,15 @@ void
 MemoryContextMemConsumed(MemoryContext context,
 						 MemoryContextCounters *consumed)
 {
+
 	memset(consumed, 0, sizeof(*consumed));
 
-	MemoryContextStatsInternal(context, 0, false, 0, consumed, false);
+	for (MemoryContext curr = context->firstchild;
+		 curr != NULL;
+		 curr = MemoryContextTraverseNext(curr, context))
+	{
+		child->methods->stats(child, NULL, NULL, consumed, print_to_stderr);
+	}
 }
 
 /*
@@ -707,8 +785,8 @@ MemoryContextMemConsumed(MemoryContext context,
 void
 MemoryContextStats(MemoryContext context)
 {
-	/* A hard-wired limit on the number of children is usually good enough */
-	MemoryContextStatsDetail(context, 100, true);
+	/* Hard-wired limits are usually good enough */
+	MemoryContextStatsDetail(context, 100, 100, true);
 }
 
 /*
@@ -720,14 +798,15 @@ MemoryContextStats(MemoryContext context)
  * with fprintf(stderr), otherwise use ereport().
  */
 void
-MemoryContextStatsDetail(MemoryContext context, int max_children,
+MemoryContextStatsDetail(MemoryContext context,
+						 int max_level, int max_children,
 						 bool print_to_stderr)
 {
 	MemoryContextCounters grand_totals;
 
 	memset(&grand_totals, 0, sizeof(grand_totals));
 
-	MemoryContextStatsInternal(context, 0, true, max_children, &grand_totals, print_to_stderr);
+	MemoryContextStatsInternal(context, 0, max_level, max_children, &grand_totals, print_to_stderr);
 
 	if (print_to_stderr)
 		fprintf(stderr,
@@ -765,11 +844,10 @@ MemoryContextStatsDetail(MemoryContext context, int max_children,
  */
 static void
 MemoryContextStatsInternal(MemoryContext context, int level,
-						   bool print, int max_children,
+						   int max_level, int max_children,
 						   MemoryContextCounters *totals,
 						   bool print_to_stderr)
 {
-	MemoryContextCounters local_totals;
 	MemoryContext child;
 	int			ichild;
 
@@ -777,65 +855,75 @@ MemoryContextStatsInternal(MemoryContext context, int level,
 
 	/* Examine the context itself */
 	context->methods->stats(context,
-							print ? MemoryContextStatsPrint : NULL,
+							MemoryContextStatsPrint,
 							(void *) &level,
-							totals, print_to_stderr);
+							totals,
+							print_to_stderr);
 
 	/*
-	 * Examine children.  If there are more than max_children of them, we do
-	 * not print the rest explicitly, but just summarize them.
+	 * Examine children.
+	 *
+	 * If we are past the recursion depth limit or already running low on
+	 * stack, do not print them explicitly but just summarize them. Similarly,
+	 * if there are more than max_children of them, we do not print the rest
+	 * explicitly, but just summarize them.
 	 */
-	memset(&local_totals, 0, sizeof(local_totals));
-
-	for (child = context->firstchild, ichild = 0;
-		 child != NULL;
-		 child = child->nextchild, ichild++)
+	ichild = 0;
+	child = context->firstchild;
+	if (level < max_level && !stack_is_too_deep())
 	{
-		if (ichild < max_children)
+		for (; child != NULL && ichild < max_children;
+			 child = child->nextchild)
+		{
 			MemoryContextStatsInternal(child, level + 1,
-									   print, max_children,
+									   max_level, max_children,
 									   totals,
 									   print_to_stderr);
-		else
-			MemoryContextStatsInternal(child, level + 1,
-									   false, max_children,
-									   &local_totals,
-									   print_to_stderr);
+			ichild++;
+		}
 	}
 
-	/* Deal with excess children */
-	if (ichild > max_children)
+	if (child != NULL)
 	{
-		if (print)
+		/* Summarize the rest of the children, avoiding recursion. */
+		MemoryContextCounters local_totals;
+
+		memset(&local_totals, 0, sizeof(local_totals));
+
+		while (child != NULL)
+		{
+			child->methods->stats(child, NULL, NULL, &local_totals, print_to_stderr);
+			child = MemoryContextTraverseNext(child, context);
+			ichild++;
+		}
+
+		if (print_to_stderr)
 		{
-			if (print_to_stderr)
-			{
-				int			i;
-
-				for (i = 0; i <= level; i++)
-					fprintf(stderr, "  ");
-				fprintf(stderr,
-						"%d more child contexts containing %zu total in %zu blocks; %zu free (%zu chunks); %zu used\n",
-						ichild - max_children,
-						local_totals.totalspace,
-						local_totals.nblocks,
-						local_totals.freespace,
-						local_totals.freechunks,
-						local_totals.totalspace - local_totals.freespace);
-			}
-			else
-				ereport(LOG_SERVER_ONLY,
-						(errhidestmt(true),
-						 errhidecontext(true),
-						 errmsg_internal("level: %d; %d more child contexts containing %zu total in %zu blocks; %zu free (%zu chunks); %zu used",
-										 level,
-										 ichild - max_children,
-										 local_totals.totalspace,
-										 local_totals.nblocks,
-										 local_totals.freespace,
-										 local_totals.freechunks,
-										 local_totals.totalspace - local_totals.freespace)));
+			int			i;
+
+			for (i = 0; i <= level; i++)
+				fprintf(stderr, "  ");
+			fprintf(stderr,
+					"%d more child contexts containing %zu total in %zu blocks; %zu free (%zu chunks); %zu used\n",
+					ichild - max_children,
+					local_totals.totalspace,
+					local_totals.nblocks,
+					local_totals.freespace,
+					local_totals.freechunks,
+					local_totals.totalspace - local_totals.freespace);
 		}
+		else
+			ereport(LOG_SERVER_ONLY,
+					(errhidestmt(true),
+					 errhidecontext(true),
+					 errmsg_internal("level: %d; %d more child contexts containing %zu total in %zu blocks; %zu free (%zu chunks); %zu used",
+									 level,
+									 ichild - max_children,
+									 local_totals.totalspace,
+									 local_totals.nblocks,
+									 local_totals.freespace,
+									 local_totals.freechunks,
+									 local_totals.totalspace - local_totals.freespace)));
 
 		if (totals)
 		{
@@ -856,8 +944,7 @@ MemoryContextStatsInternal(MemoryContext context, int level,
  */
 static void
 MemoryContextStatsPrint(MemoryContext context, void *passthru,
-						const char *stats_string,
-						bool print_to_stderr)
+						const char *stats_string, bool print_to_stderr)
 {
 	int			level = *(int *) passthru;
 	const char *name = context->name;
@@ -936,13 +1023,15 @@ MemoryContextStatsPrint(MemoryContext context, void *passthru,
 void
 MemoryContextCheck(MemoryContext context)
 {
-	MemoryContext child;
-
 	Assert(MemoryContextIsValid(context));
-
 	context->methods->check(context);
-	for (child = context->firstchild; child != NULL; child = child->nextchild)
-		MemoryContextCheck(child);
+
+	for (MemoryContext curr = context->firstchild;
+		 curr != NULL;
+		 curr = MemoryContextTraverseNext(curr, context))
+	{
+		curr->methods->check(curr);
+	}
 }
 #endif
 
@@ -1183,14 +1272,15 @@ ProcessLogMemoryContextInterrupt(void)
 	/*
 	 * When a backend process is consuming huge memory, logging all its memory
 	 * contexts might overrun available disk space. To prevent this, we limit
-	 * the number of child contexts to log per parent to 100.
+	 * the depth of the hierarchy, as well as the number of child contexts to
+	 * log per parent to 100.
 	 *
 	 * As with MemoryContextStats(), we suppose that practical cases where the
 	 * dump gets long will typically be huge numbers of siblings under the
 	 * same parent context; while the additional debugging value from seeing
 	 * details about individual siblings beyond 100 will not be large.
 	 */
-	MemoryContextStatsDetail(TopMemoryContext, 100, false);
+	MemoryContextStatsDetail(TopMemoryContext, 100, 100, false);
 }
 
 void *
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index 7fd41d20caf..6e5fa72b0e1 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -87,7 +87,8 @@ extern Size MemoryContextMemAllocated(MemoryContext context, bool recurse);
 extern void MemoryContextMemConsumed(MemoryContext context,
 									 MemoryContextCounters *consumed);
 extern void MemoryContextStats(MemoryContext context);
-extern void MemoryContextStatsDetail(MemoryContext context, int max_children,
+extern void MemoryContextStatsDetail(MemoryContext context,
+									 int max_level, int max_children,
 									 bool print_to_stderr);
 extern void MemoryContextAllowInCriticalSection(MemoryContext context,
 												bool allow);
-- 
2.39.3 (Apple Git-145)

