From e4713160a80e95cc99b043af11f8c33f500b5786 Mon Sep 17 00:00:00 2001
From: Melih Mutlu <m.melihmutlu@gmail.com>
Date: Sat, 13 Jul 2024 00:59:26 +0300
Subject: [PATCH v10] Add path column into pg_backend_memory_contexts

This patch adds a new column into the tuples returned by
pg_get_backend_memory_contexts() to specify the parent/child relation
between memory contexts and the path from root to current context.

Context names cannot be relied on since they're not unique. Therefore,
unique IDs are needed for each context. Those new IDs are assigned during
pg_get_backend_memory_contexts() call and not stored anywhere. So they
may change in each pg_get_backend_memory_contexts() call and shouldn't be
used across different pg_get_backend_memory_contexts() calls.
---
 doc/src/sgml/system-views.sgml         |  40 ++++-
 src/backend/utils/adt/mcxtfuncs.c      | 195 ++++++++++++++++++++-----
 src/include/catalog/pg_proc.dat        |   6 +-
 src/include/nodes/memnodes.h           |   4 +-
 src/test/regress/expected/rules.out    |   3 +-
 src/test/regress/expected/sysviews.out |  18 ++-
 src/test/regress/sql/sysviews.sql      |  12 +-
 7 files changed, 232 insertions(+), 46 deletions(-)

diff --git a/doc/src/sgml/system-views.sgml b/doc/src/sgml/system-views.sgml
index bdc34cf94e..9602894dae 100644
--- a/doc/src/sgml/system-views.sgml
+++ b/doc/src/sgml/system-views.sgml
@@ -504,7 +504,25 @@
        <structfield>level</structfield> <type>int4</type>
       </para>
       <para>
-       Distance from TopMemoryContext in context tree
+       Represents the level in the memory context hierarchy tree. Level of
+       a context also shows the position of that context in
+       <structfield>path</structfield> arrays. TopMemoryContext has the lowest
+       level which is 1.
+      </para></entry>
+     </row>
+
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>path</structfield> <type>int4[]</type>
+      </para>
+      <para>
+       Array of transient identifiers to describe the memory context hierarchy.
+       A path includes all contexts from TopMemoryContext to the current
+       context in that hierarchy. The first array element is always the
+       TopMemoryContext and the last element in the list refers to the current
+       context. Note that these IDs are unstable between multiple invocations
+       of the view. See the example query below for advice on how to use this
+       column effectively
       </para></entry>
      </row>
 
@@ -561,6 +579,26 @@
    read only by superusers or roles with the privileges of the
    <literal>pg_read_all_stats</literal> role.
   </para>
+
+  <para>
+   The <structfield>path</structfield> column can be useful to build
+   parent/child relation between memory contexts. For example, the following
+   query calculates the total number of bytes used by CacheMemoryContext and
+   its child contexts:
+<programlisting>
+WITH memory_contexts AS (
+    SELECT *
+    FROM pg_backend_memory_contexts
+)
+SELECT sum(c1.total_bytes)
+FROM memory_contexts c1, memory_contexts c2
+WHERE c2.name = 'CacheMemoryContext'
+AND c1.path[c2.level] = c2.path[c2.level];
+</programlisting>
+    Also, <link linkend="queries-with">Common Table Expressions</link> can be
+    useful while working with context IDs as these IDs are temporary and may
+    change in each invocation.
+  </para>
  </sect1>
 
  <sect1 id="view-pg-config">
diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
index 1085941484..2235457997 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -17,9 +17,12 @@
 
 #include "funcapi.h"
 #include "mb/pg_wchar.h"
+#include "miscadmin.h"
 #include "storage/proc.h"
 #include "storage/procarray.h"
+#include "utils/array.h"
 #include "utils/builtins.h"
+#include "utils/hsearch.h"
 
 /* ----------
  * The max bytes for showing identifiers of MemoryContext.
@@ -27,47 +30,103 @@
  */
 #define MEMORY_CONTEXT_IDENT_DISPLAY_SIZE	1024
 
+typedef struct MemoryContextId
+{
+	MemoryContext context;
+	int context_id;
+} MemoryContextId;
+
+/*
+ * get_memory_context_name_and_ident
+ *		Populate *name and *ident from the name and ident from 'context'.
+ */
+static void
+get_memory_context_name_and_ident(MemoryContext context, const char **name,
+								   const char **ident)
+{
+	*name = context->name;
+	*ident = context->ident;
+
+	/*
+	* To be consistent with logging output, we label dynahash contexts with
+	* just the hash table name as with MemoryContextStatsPrint().
+	*/
+	if (ident && strcmp(*name, "dynahash") == 0)
+	{
+		*name = *ident;
+		*ident = NULL;
+	}
+}
+
+/*
+ * int_list_to_array
+ *		Convert a IntList to an int[] array.
+ */
+static Datum
+int_list_to_array(List *list)
+{
+	Datum *datum_array;
+	int length;
+	ArrayType *result_array;
+
+	length = list_length(list);
+	datum_array = (Datum *) palloc(length * sizeof(Datum));
+	length = 0;
+	foreach_int(id, list)
+		datum_array[length++] = Int32GetDatum(id);
+
+	result_array = construct_array_builtin(datum_array, length, INT4OID);
+
+	return PointerGetDatum(result_array);
+}
+
 /*
  * PutMemoryContextsStatsTupleStore
- *		One recursion level for pg_get_backend_memory_contexts.
+ *		Add details for the given MemoryContext to 'tupstore'.
  */
 static void
 PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
-								 TupleDesc tupdesc, MemoryContext context,
-								 const char *parent, int level)
+								 TupleDesc tupdesc, MemoryContext context, HTAB *context_id_lookup)
 {
-#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	10
+#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	11
 
 	Datum		values[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
 	bool		nulls[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
 	MemoryContextCounters stat;
-	MemoryContext child;
+	List *path = NIL;
 	const char *name;
 	const char *ident;
 	const char *type;
 
 	Assert(MemoryContextIsValid(context));
 
-	name = context->name;
-	ident = context->ident;
-
 	/*
-	 * To be consistent with logging output, we label dynahash contexts with
-	 * just the hash table name as with MemoryContextStatsPrint().
+	 * Figure out the transient context_id of this context and each of its
+	 * ancestors.
 	 */
-	if (ident && strcmp(name, "dynahash") == 0)
+	for (MemoryContext cur = context; cur != NULL; cur = cur->parent)
 	{
-		name = ident;
-		ident = NULL;
+		MemoryContextId *entry;
+		bool	found;
+
+		entry = hash_search(context_id_lookup, &cur, HASH_FIND, &found);
+
+		if (!found)
+			elog(ERROR, "hash table corrupted");
+		path = lcons_int(entry->context_id, path);
 	}
 
 	/* Examine the context itself */
 	memset(&stat, 0, sizeof(stat));
-	(*context->methods->stats) (context, NULL, (void *) &level, &stat, true);
+	(*context->methods->stats)(context, NULL, NULL, &stat, true);
 
 	memset(values, 0, sizeof(values));
 	memset(nulls, 0, sizeof(nulls));
 
+	get_memory_context_name_and_ident(context,
+									  (const char **)&name,
+									  (const char **) &ident);
+
 	if (name)
 		values[0] = CStringGetTextDatum(name);
 	else
@@ -75,15 +134,17 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 
 	if (ident)
 	{
-		int			idlen = strlen(ident);
-		char		clipped_ident[MEMORY_CONTEXT_IDENT_DISPLAY_SIZE];
+		int idlen = strlen(ident);
+		char clipped_ident[MEMORY_CONTEXT_IDENT_DISPLAY_SIZE];
 
 		/*
-		 * Some identifiers such as SQL query string can be very long,
-		 * truncate oversize identifiers.
-		 */
+		* Some identifiers such as SQL query string can be very long,
+		* truncate oversize identifiers.
+		*/
 		if (idlen >= MEMORY_CONTEXT_IDENT_DISPLAY_SIZE)
-			idlen = pg_mbcliplen(ident, idlen, MEMORY_CONTEXT_IDENT_DISPLAY_SIZE - 1);
+			idlen = pg_mbcliplen(ident,
+									idlen,
+									MEMORY_CONTEXT_IDENT_DISPLAY_SIZE - 1);
 
 		memcpy(clipped_ident, ident, idlen);
 		clipped_ident[idlen] = '\0';
@@ -92,8 +153,15 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	else
 		nulls[1] = true;
 
-	if (parent)
-		values[2] = CStringGetTextDatum(parent);
+	if (context->parent)
+	{
+		char *parent_name, *parent_ident;
+
+		get_memory_context_name_and_ident(context->parent,
+										  (const char **)&parent_name,
+										  (const char **)&parent_ident);
+		values[2] = CStringGetTextDatum(parent_name);
+	}
 	else
 		nulls[2] = true;
 
@@ -117,19 +185,16 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	}
 
 	values[3] = CStringGetTextDatum(type);
-	values[4] = Int32GetDatum(level);
-	values[5] = Int64GetDatum(stat.totalspace);
-	values[6] = Int64GetDatum(stat.nblocks);
-	values[7] = Int64GetDatum(stat.freespace);
-	values[8] = Int64GetDatum(stat.freechunks);
-	values[9] = Int64GetDatum(stat.totalspace - stat.freespace);
-	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+	values[4] = Int32GetDatum(list_length(path)); /* level */
+	values[5] = int_list_to_array(path);
+	values[6] = Int64GetDatum(stat.totalspace);
+	values[7] = Int64GetDatum(stat.nblocks);
+	values[8] = Int64GetDatum(stat.freespace);
+	values[9] = Int64GetDatum(stat.freechunks);
+	values[10] = Int64GetDatum(stat.totalspace - stat.freespace);
 
-	for (child = context->firstchild; child != NULL; child = child->nextchild)
-	{
-		PutMemoryContextsStatsTupleStore(tupstore, tupdesc,
-										 child, name, level + 1);
-	}
+	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+	list_free(path);
 }
 
 /*
@@ -140,10 +205,68 @@ Datum
 pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 {
 	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	int context_id;
+	List *queue;
+	ListCell *lc;
+	HASHCTL ctl;
+	HTAB *context_id_lookup;
+
+
+	ctl.keysize = sizeof(MemoryContext);
+	ctl.entrysize = sizeof(MemoryContextId);
+
+	context_id_lookup = hash_create("pg_get_backend_memory_contexts lookup",
+									256,
+									&ctl,
+									HASH_ELEM | HASH_BLOBS);
 
 	InitMaterializedSRF(fcinfo, 0);
-	PutMemoryContextsStatsTupleStore(rsinfo->setResult, rsinfo->setDesc,
-									 TopMemoryContext, NULL, 0);
+
+	/*
+	 * Here we use a non-recursive algorithm to visit all MemoryContexts
+	 * starting with TopMemoryContext.  The reason we avoid using a recursive
+	 * algorithm is because we want to assign the context_id breadth first.
+	 * I.e. all context at level 1 are assigned ids before contexts at level 2.
+	 * Because lower-leveled contexts are less likely to change, this makes the
+	 * assigned context_id more stable.  Otherwise, if the first child of
+	 * TopMemoryContext obtained an additional grand child, the context_id for
+	 * the second child of TopMemoryContext would change.
+	 */
+	queue = list_make1(TopMemoryContext);
+
+	/* TopMemoryContext will always have a context_id of 1 */
+	context_id = 1;
+
+	foreach(lc, queue)
+	{
+		MemoryContext cur = lfirst(lc);
+		MemoryContextId *entry;
+		bool found;
+
+		/*
+		 * Record the context_id that we've assigned to each MemoryContext.
+		 * PutMemoryContextsStatsTupleStore needs this to populate the "path"
+		 * column with the parent context_ids.
+		 */
+		entry = (MemoryContextId *) hash_search(context_id_lookup, &cur,
+												HASH_ENTER, &found);
+		entry->context_id = context_id++;
+		Assert(!found);
+
+		PutMemoryContextsStatsTupleStore(rsinfo->setResult,
+										 rsinfo->setDesc,
+										 cur,
+										 context_id_lookup);
+
+		/*
+		 * Queue up all the child contexts of this level for the next
+		 * iteration.
+		 */
+		for (MemoryContext c = cur->firstchild; c != NULL; c = c->nextchild)
+			queue = lappend(queue, c);
+	}
+
+	hash_destroy(context_id_lookup);
 
 	return (Datum) 0;
 }
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 73d9cf8582..d14a94b987 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8290,9 +8290,9 @@
   proname => 'pg_get_backend_memory_contexts', prorows => '100',
   proretset => 't', provolatile => 'v', proparallel => 'r',
   prorettype => 'record', proargtypes => '',
-  proallargtypes => '{text,text,text,text,int4,int8,int8,int8,int8,int8}',
-  proargmodes => '{o,o,o,o,o,o,o,o,o,o}',
-  proargnames => '{name, ident, parent, type, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes}',
+  proallargtypes => '{text,text,text,text,int4,_int4,int8,int8,int8,int8,int8}',
+  proargmodes => '{o,o,o,o,o,o,o,o,o,o,o}',
+  proargnames => '{name, ident, parent, type, level, path, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes}',
   prosrc => 'pg_get_backend_memory_contexts' },
 
 # logging memory contexts of the specified backend
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index c4c9fd3e3e..addfbb1ccd 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -128,8 +128,8 @@ typedef struct MemoryContextData
 	MemoryContext firstchild;	/* head of linked list of children */
 	MemoryContext prevchild;	/* previous child of same parent */
 	MemoryContext nextchild;	/* next child of same parent */
-	const char *name;			/* context name (just for debugging) */
-	const char *ident;			/* context ID if any (just for debugging) */
+	const char *name;			/* context name */
+	const char *ident;			/* context ID if any */
 	MemoryContextCallback *reset_cbs;	/* list of reset/delete callbacks */
 } MemoryContextData;
 
diff --git a/src/test/regress/expected/rules.out b/src/test/regress/expected/rules.out
index 4c789279e5..5201280669 100644
--- a/src/test/regress/expected/rules.out
+++ b/src/test/regress/expected/rules.out
@@ -1308,12 +1308,13 @@ pg_backend_memory_contexts| SELECT name,
     parent,
     type,
     level,
+    path,
     total_bytes,
     total_nblocks,
     free_bytes,
     free_chunks,
     used_bytes
-   FROM pg_get_backend_memory_contexts() pg_get_backend_memory_contexts(name, ident, parent, type, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes);
+   FROM pg_get_backend_memory_contexts() pg_get_backend_memory_contexts(name, ident, parent, type, level, path, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes);
 pg_config| SELECT name,
     setting
    FROM pg_config() pg_config(name, setting);
diff --git a/src/test/regress/expected/sysviews.out b/src/test/regress/expected/sysviews.out
index 729620de13..7718c06835 100644
--- a/src/test/regress/expected/sysviews.out
+++ b/src/test/regress/expected/sysviews.out
@@ -22,10 +22,10 @@ select count(*) >= 0 as ok from pg_available_extensions;
 -- The entire output of pg_backend_memory_contexts is not stable,
 -- we test only the existence and basic condition of TopMemoryContext.
 select type, name, ident, parent, level, total_bytes >= free_bytes
-  from pg_backend_memory_contexts where level = 0;
+  from pg_backend_memory_contexts where level = 1;
    type   |       name       | ident | parent | level | ?column? 
 ----------+------------------+-------+--------+-------+----------
- AllocSet | TopMemoryContext |       |        |     0 | t
+ AllocSet | TopMemoryContext |       |        |     1 | t
 (1 row)
 
 -- We can exercise some MemoryContext type stats functions.  Most of the
@@ -51,6 +51,20 @@ from pg_backend_memory_contexts where name = 'Caller tuples';
 (1 row)
 
 rollback;
+-- Test whether there are contexts with CacheMemoryContext in their path.
+-- There should be multiple children of CacheMemoryContext.
+with contexts as (
+  select * from pg_backend_memory_contexts
+)
+select count(*) > 1
+from contexts c1, contexts c2
+where c2.name = 'CacheMemoryContext'
+and c1.path[c2.level] = c2.path[c2.level];
+ ?column? 
+----------
+ t
+(1 row)
+
 -- At introduction, pg_config had 23 entries; it may grow
 select count(*) > 20 as ok from pg_config;
  ok 
diff --git a/src/test/regress/sql/sysviews.sql b/src/test/regress/sql/sysviews.sql
index 7edac2fde1..b3f9c55c33 100644
--- a/src/test/regress/sql/sysviews.sql
+++ b/src/test/regress/sql/sysviews.sql
@@ -15,7 +15,7 @@ select count(*) >= 0 as ok from pg_available_extensions;
 -- The entire output of pg_backend_memory_contexts is not stable,
 -- we test only the existence and basic condition of TopMemoryContext.
 select type, name, ident, parent, level, total_bytes >= free_bytes
-  from pg_backend_memory_contexts where level = 0;
+  from pg_backend_memory_contexts where level = 1;
 
 -- We can exercise some MemoryContext type stats functions.  Most of the
 -- column values are too platform-dependant to display.
@@ -32,6 +32,16 @@ select type, name, parent, total_bytes > 0, total_nblocks, free_bytes > 0, free_
 from pg_backend_memory_contexts where name = 'Caller tuples';
 rollback;
 
+-- Test whether there are contexts with CacheMemoryContext in their path.
+-- There should be multiple children of CacheMemoryContext.
+with contexts as (
+  select * from pg_backend_memory_contexts
+)
+select count(*) > 1
+from contexts c1, contexts c2
+where c2.name = 'CacheMemoryContext'
+and c1.path[c2.level] = c2.path[c2.level];
+
 -- At introduction, pg_config had 23 entries; it may grow
 select count(*) > 20 as ok from pg_config;
 
-- 
2.34.1

