Parent/child context relation in pg_get_backend_memory_contexts()

From e9134ad5465d0dae626cf658def2ca58648d64c2 Mon Sep 17 00:00:00 2001
From: Melih Mutlu <m.melihmutlu@gmail.com>
Date: Tue, 13 Jun 2023 15:00:05 +0300
Subject: [PATCH v2] Adding id/parent_id into pg_backend_memory_contexts

Added three new fields 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 id and parent_id are needed for each context. Those new id's 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 differenct
pg_get_backend_memory_contexts() calls.

Context id's are start from 0 which is assigned to TopMemoryContext.
---
 doc/src/sgml/system-views.sgml      | 27 +++++++++++++
 src/backend/utils/adt/mcxtfuncs.c   | 59 ++++++++++++++++++++++++++---
 src/include/catalog/pg_proc.dat     |  6 +--
 src/test/regress/expected/rules.out |  7 +++-
 4 files changed, 89 insertions(+), 10 deletions(-)

diff --git a/doc/src/sgml/system-views.sgml b/doc/src/sgml/system-views.sgml
index 57b228076e..7811130cfb 100644
--- a/doc/src/sgml/system-views.sgml
+++ b/doc/src/sgml/system-views.sgml
@@ -538,6 +538,33 @@
        Used space in bytes
       </para></entry>
      </row>
+
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>context_id</structfield> <type>int4</type>
+      </para>
+      <para>
+       Current context id
+      </para></entry>
+     </row>
+
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>parent_id</structfield> <type>int4</type>
+      </para>
+      <para>
+       Parent context id
+      </para></entry>
+     </row>
+
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>path</structfield> <type>int4</type>
+      </para>
+      <para>
+       Path to reach the current context from TopMemoryContext
+      </para></entry>
+     </row>
     </tbody>
    </tgroup>
   </table>
diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
index 92ca5b2f72..81cb35dd47 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -20,6 +20,7 @@
 #include "mb/pg_wchar.h"
 #include "storage/proc.h"
 #include "storage/procarray.h"
+#include "utils/array.h"
 #include "utils/builtins.h"
 
 /* ----------
@@ -28,6 +29,8 @@
  */
 #define MEMORY_CONTEXT_IDENT_DISPLAY_SIZE	1024
 
+static Datum convert_path_to_datum(List *path);
+
 /*
  * PutMemoryContextsStatsTupleStore
  *		One recursion level for pg_get_backend_memory_contexts.
@@ -35,9 +38,10 @@
 static void
 PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 								 TupleDesc tupdesc, MemoryContext context,
-								 const char *parent, int level)
+								 const char *parent, int level, int *context_id,
+								 int parent_id, List *path)
 {
-#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	9
+#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	12
 
 	Datum		values[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
 	bool		nulls[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
@@ -45,6 +49,7 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	MemoryContext child;
 	const char *name;
 	const char *ident;
+	int  current_context_id = (*context_id)++;
 
 	Assert(MemoryContextIsValid(context));
 
@@ -103,13 +108,29 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	values[6] = Int64GetDatum(stat.freespace);
 	values[7] = Int64GetDatum(stat.freechunks);
 	values[8] = Int64GetDatum(stat.totalspace - stat.freespace);
+	values[9] = Int32GetDatum(current_context_id);
+
+	if(parent_id < 0)
+		/* TopMemoryContext has no parent context */
+		nulls[10] = true;
+	else
+		values[10] = Int32GetDatum(parent_id);
+
+	if (path == NIL)
+		nulls[11] = true;
+	else
+		values[11] = convert_path_to_datum(path);
+
 	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
 
+	path = lappend_int(path, current_context_id);
 	for (child = context->firstchild; child != NULL; child = child->nextchild)
 	{
-		PutMemoryContextsStatsTupleStore(tupstore, tupdesc,
-										 child, name, level + 1);
+		PutMemoryContextsStatsTupleStore(tupstore, tupdesc, child, name,
+										 level+1, context_id,
+										 current_context_id, path);
 	}
+	path = list_delete_last(path);
 }
 
 /*
@@ -120,10 +141,15 @@ Datum
 pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 {
 	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	int context_id = 0;
+	List *path = NIL;
+
+	elog(LOG, "pg_get_backend_memory_contexts called");
 
 	InitMaterializedSRF(fcinfo, 0);
 	PutMemoryContextsStatsTupleStore(rsinfo->setResult, rsinfo->setDesc,
-									 TopMemoryContext, NULL, 0);
+									 TopMemoryContext, NULL, 0, &context_id,
+									 -1, path);
 
 	return (Datum) 0;
 }
@@ -193,3 +219,26 @@ pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
 
 	PG_RETURN_BOOL(true);
 }
+
+/*
+ * Convert a list of context ids to a int[] Datum
+ */
+static Datum
+convert_path_to_datum(List *path)
+{
+	Datum	   *datum_array;
+	int			length;
+	ArrayType  *result_array;
+	ListCell   *lc;
+
+	length = list_length(path);
+	datum_array = (Datum *) palloc(length * sizeof(Datum));
+	length = 0;
+	foreach(lc, path)
+	{
+		datum_array[length++] = Int32GetDatum((int) lfirst_int(lc));
+	}
+	result_array = construct_array_builtin(datum_array, length, INT4OID);
+
+	return PointerGetDatum(result_array);
+}
\ No newline at end of file
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 6996073989..49514ad792 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8212,9 +8212,9 @@
   proname => 'pg_get_backend_memory_contexts', prorows => '100',
   proretset => 't', provolatile => 'v', proparallel => 'r',
   prorettype => 'record', proargtypes => '',
-  proallargtypes => '{text,text,text,int4,int8,int8,int8,int8,int8}',
-  proargmodes => '{o,o,o,o,o,o,o,o,o}',
-  proargnames => '{name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes}',
+  proallargtypes => '{text,text,text,int4,int8,int8,int8,int8,int8,int4,int4,_int4}',
+  proargmodes => '{o,o,o,o,o,o,o,o,o,o,o,o}',
+  proargnames => '{name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes, id, parent_id, path}',
   prosrc => 'pg_get_backend_memory_contexts' },
 
 # logging memory contexts of the specified backend
diff --git a/src/test/regress/expected/rules.out b/src/test/regress/expected/rules.out
index e07afcd4aa..c7112be09e 100644
--- a/src/test/regress/expected/rules.out
+++ b/src/test/regress/expected/rules.out
@@ -1311,8 +1311,11 @@ pg_backend_memory_contexts| SELECT name,
     total_nblocks,
     free_bytes,
     free_chunks,
-    used_bytes
-   FROM pg_get_backend_memory_contexts() pg_get_backend_memory_contexts(name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes);
+    used_bytes,
+    id,
+    parent_id,
+    path
+   FROM pg_get_backend_memory_contexts() pg_get_backend_memory_contexts(name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes, id, parent_id, path);
 pg_config| SELECT name,
     setting
    FROM pg_config() pg_config(name, setting);
-- 
2.25.1

From 4595e8e31f7c234426e7f104890cd3f0f0e4ca10 Mon Sep 17 00:00:00 2001
From: Melih Mutlu <m.melihmutlu@gmail.com>
Date: Tue, 13 Jun 2023 15:00:05 +0300
Subject: [PATCH v3] Adding id/parent_id into pg_backend_memory_contexts

Added three new fields 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 id and parent_id are needed for each context. Those new id's 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 differenct
pg_get_backend_memory_contexts() calls.

Context id's are start from 0 which is assigned to TopMemoryContext.
---
 doc/src/sgml/system-views.sgml         | 30 +++++++++++++
 src/backend/utils/adt/mcxtfuncs.c      | 62 +++++++++++++++++++++++---
 src/include/catalog/pg_proc.dat        |  6 +--
 src/test/regress/expected/rules.out    |  7 ++-
 src/test/regress/expected/sysviews.out | 23 +++++++++-
 src/test/regress/sql/sysviews.sql      | 15 ++++++-
 6 files changed, 130 insertions(+), 13 deletions(-)

diff --git a/doc/src/sgml/system-views.sgml b/doc/src/sgml/system-views.sgml
index 2b35c2f91b..29db104c81 100644
--- a/doc/src/sgml/system-views.sgml
+++ b/doc/src/sgml/system-views.sgml
@@ -543,6 +543,36 @@
        Used space in bytes
       </para></entry>
      </row>
+
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>context_id</structfield> <type>int4</type>
+      </para>
+      <para>
+       Current context id. Note that the context id is a temporary id and may
+       change in each invocation
+      </para></entry>
+     </row>
+
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>path</structfield> <type>int4</type>
+      </para>
+      <para>
+       Path to reach the current context from TopMemoryContext. Context ids in
+       this list represents all parents of the current context. This can be
+       used to build the parent and child relation.
+      </para></entry>
+     </row>
+
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>total_bytes_including_children</structfield> <type>int8</type>
+      </para>
+      <para>
+       Total bytes allocated for this memory context including its children
+      </para></entry>
+     </row>
     </tbody>
    </tgroup>
   </table>
diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
index 92ca5b2f72..9e7aa3535b 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -20,6 +20,7 @@
 #include "mb/pg_wchar.h"
 #include "storage/proc.h"
 #include "storage/procarray.h"
+#include "utils/array.h"
 #include "utils/builtins.h"
 
 /* ----------
@@ -28,6 +29,8 @@
  */
 #define MEMORY_CONTEXT_IDENT_DISPLAY_SIZE	1024
 
+static Datum convert_path_to_datum(List *path);
+
 /*
  * PutMemoryContextsStatsTupleStore
  *		One recursion level for pg_get_backend_memory_contexts.
@@ -35,9 +38,10 @@
 static void
 PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 								 TupleDesc tupdesc, MemoryContext context,
-								 const char *parent, int level)
+								 const char *parent, int level, int *context_id,
+								 List *path, Size *total_bytes_inc_chidlren)
 {
-#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	9
+#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	12
 
 	Datum		values[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
 	bool		nulls[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
@@ -45,6 +49,8 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	MemoryContext child;
 	const char *name;
 	const char *ident;
+	int  current_context_id = (*context_id)++;
+	Size total_bytes_children = 0;
 
 	Assert(MemoryContextIsValid(context));
 
@@ -103,13 +109,28 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	values[6] = Int64GetDatum(stat.freespace);
 	values[7] = Int64GetDatum(stat.freechunks);
 	values[8] = Int64GetDatum(stat.totalspace - stat.freespace);
-	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+	values[9] = Int32GetDatum(current_context_id);
+
+	if (path == NIL)
+		nulls[10] = true;
+	else
+		values[10] = convert_path_to_datum(path);
 
+	path = lappend_int(path, current_context_id);
 	for (child = context->firstchild; child != NULL; child = child->nextchild)
 	{
-		PutMemoryContextsStatsTupleStore(tupstore, tupdesc,
-										 child, name, level + 1);
+		PutMemoryContextsStatsTupleStore(tupstore, tupdesc, child, name,
+										 level+1, context_id, path,
+										 &total_bytes_children);
 	}
+	path = list_delete_last(path);
+
+	total_bytes_children += stat.totalspace;
+	values[11] = Int64GetDatum(total_bytes_children);
+
+	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+
+	*total_bytes_inc_chidlren += total_bytes_children;
 }
 
 /*
@@ -120,10 +141,16 @@ Datum
 pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 {
 	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	int context_id = 0;
+	List *path = NIL;
+	Size total_bytes_inc_chidlren = 0;
+
+	elog(LOG, "pg_get_backend_memory_contexts called");
 
 	InitMaterializedSRF(fcinfo, 0);
 	PutMemoryContextsStatsTupleStore(rsinfo->setResult, rsinfo->setDesc,
-									 TopMemoryContext, NULL, 0);
+									 TopMemoryContext, NULL, 0, &context_id,
+									 path, &total_bytes_inc_chidlren);
 
 	return (Datum) 0;
 }
@@ -193,3 +220,26 @@ pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
 
 	PG_RETURN_BOOL(true);
 }
+
+/*
+ * Convert a list of context ids to a int[] Datum
+ */
+static Datum
+convert_path_to_datum(List *path)
+{
+	Datum	   *datum_array;
+	int			length;
+	ArrayType  *result_array;
+	ListCell   *lc;
+
+	length = list_length(path);
+	datum_array = (Datum *) palloc(length * sizeof(Datum));
+	length = 0;
+	foreach(lc, path)
+	{
+		datum_array[length++] = Int32GetDatum(lfirst_int(lc));
+	}
+	result_array = construct_array_builtin(datum_array, length, INT4OID);
+
+	return PointerGetDatum(result_array);
+}
\ No newline at end of file
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index c92d0631a0..789ba25fa3 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8239,9 +8239,9 @@
   proname => 'pg_get_backend_memory_contexts', prorows => '100',
   proretset => 't', provolatile => 'v', proparallel => 'r',
   prorettype => 'record', proargtypes => '',
-  proallargtypes => '{text,text,text,int4,int8,int8,int8,int8,int8}',
-  proargmodes => '{o,o,o,o,o,o,o,o,o}',
-  proargnames => '{name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes}',
+  proallargtypes => '{text,text,text,int4,int8,int8,int8,int8,int8,int4,_int4,int8}',
+  proargmodes => '{o,o,o,o,o,o,o,o,o,o,o,o}',
+  proargnames => '{name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes, id, path, total_bytes_including_children}',
   prosrc => 'pg_get_backend_memory_contexts' },
 
 # logging memory contexts of the specified backend
diff --git a/src/test/regress/expected/rules.out b/src/test/regress/expected/rules.out
index 2c60400ade..8fd3e9f6fd 100644
--- a/src/test/regress/expected/rules.out
+++ b/src/test/regress/expected/rules.out
@@ -1311,8 +1311,11 @@ pg_backend_memory_contexts| SELECT name,
     total_nblocks,
     free_bytes,
     free_chunks,
-    used_bytes
-   FROM pg_get_backend_memory_contexts() pg_get_backend_memory_contexts(name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes);
+    used_bytes,
+    id,
+    path,
+    total_bytes_including_children
+   FROM pg_get_backend_memory_contexts() pg_get_backend_memory_contexts(name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes, id, path, total_bytes_including_children);
 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 aae5d51e1c..45d8064c35 100644
--- a/src/test/regress/expected/sysviews.out
+++ b/src/test/regress/expected/sysviews.out
@@ -20,7 +20,7 @@ select count(*) >= 0 as ok from pg_available_extensions;
 (1 row)
 
 -- The entire output of pg_backend_memory_contexts is not stable,
--- we test only the existence and basic condition of TopMemoryContext.
+-- we test the existence and basic condition of TopMemoryContext.
 select name, ident, parent, level, total_bytes >= free_bytes
   from pg_backend_memory_contexts where level = 0;
        name       | ident | parent | level | ?column? 
@@ -28,6 +28,27 @@ select name, ident, parent, level, total_bytes >= free_bytes
  TopMemoryContext |       |        |     0 | t
 (1 row)
 
+-- 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(*) > 0 
+from contexts 
+where array[(select id from contexts where name = 'CacheMemoryContext')] <@ path;
+ ?column? 
+----------
+ t
+(1 row)
+
+-- TopMemoryContext should have the largest total_bytes_including_children.
+select name from pg_backend_memory_contexts
+  order by total_bytes_including_children desc limit 1;
+       name       
+------------------
+ TopMemoryContext
+(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 6b4e24601d..56cfcb77f4 100644
--- a/src/test/regress/sql/sysviews.sql
+++ b/src/test/regress/sql/sysviews.sql
@@ -13,10 +13,23 @@ select count(*) >= 0 as ok from pg_available_extension_versions;
 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.
+-- we test the existence and basic condition of TopMemoryContext.
 select name, ident, parent, level, total_bytes >= free_bytes
   from pg_backend_memory_contexts where level = 0;
 
+-- 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(*) > 0 
+from contexts 
+where array[(select id from contexts where name = 'CacheMemoryContext')] <@ path;
+
+-- TopMemoryContext should have the largest total_bytes_including_children.
+select name from pg_backend_memory_contexts
+  order by total_bytes_including_children desc limit 1;
+
 -- At introduction, pg_config had 23 entries; it may grow
 select count(*) > 20 as ok from pg_config;
 
-- 
2.34.1

From 8971207cc727870f4fbaad76579f901247fb0c1d Mon Sep 17 00:00:00 2001
From: Melih Mutlu <m.melihmutlu@gmail.com>
Date: Tue, 13 Jun 2023 15:00:05 +0300
Subject: [PATCH v4] Adding id/parent_id into pg_backend_memory_contexts

Added three new fields 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 id and parent_id are needed for each context. Those new id's 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 differenct
pg_get_backend_memory_contexts() calls.

Context id's are start from 0 which is assigned to TopMemoryContext.
---
 doc/src/sgml/system-views.sgml         | 30 +++++++++++++
 src/backend/utils/adt/mcxtfuncs.c      | 60 +++++++++++++++++++++++---
 src/include/catalog/pg_proc.dat        |  6 +--
 src/test/regress/expected/rules.out    |  7 ++-
 src/test/regress/expected/sysviews.out | 23 +++++++++-
 src/test/regress/sql/sysviews.sql      | 15 ++++++-
 6 files changed, 128 insertions(+), 13 deletions(-)

diff --git a/doc/src/sgml/system-views.sgml b/doc/src/sgml/system-views.sgml
index 0ef1745631..488e9df4f6 100644
--- a/doc/src/sgml/system-views.sgml
+++ b/doc/src/sgml/system-views.sgml
@@ -543,6 +543,36 @@
        Used space in bytes
       </para></entry>
      </row>
+
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>context_id</structfield> <type>int4</type>
+      </para>
+      <para>
+       Current context id. Note that the context id is a temporary id and may
+       change in each invocation
+      </para></entry>
+     </row>
+
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>path</structfield> <type>int4[]</type>
+      </para>
+      <para>
+       Path to reach the current context from TopMemoryContext. Context ids in
+       this list represents all parents of the current context. This can be
+       used to build the parent and child relation
+      </para></entry>
+     </row>
+
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>total_bytes_including_children</structfield> <type>int8</type>
+      </para>
+      <para>
+       Total bytes allocated for this memory context including its children
+      </para></entry>
+     </row>
     </tbody>
    </tgroup>
   </table>
diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
index 92ca5b2f72..bc357a083b 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -20,6 +20,7 @@
 #include "mb/pg_wchar.h"
 #include "storage/proc.h"
 #include "storage/procarray.h"
+#include "utils/array.h"
 #include "utils/builtins.h"
 
 /* ----------
@@ -28,6 +29,8 @@
  */
 #define MEMORY_CONTEXT_IDENT_DISPLAY_SIZE	1024
 
+static Datum convert_path_to_datum(List *path);
+
 /*
  * PutMemoryContextsStatsTupleStore
  *		One recursion level for pg_get_backend_memory_contexts.
@@ -35,9 +38,10 @@
 static void
 PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 								 TupleDesc tupdesc, MemoryContext context,
-								 const char *parent, int level)
+								 const char *parent, int level, int *context_id,
+								 List *path, Size *total_bytes_inc_children)
 {
-#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	9
+#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	12
 
 	Datum		values[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
 	bool		nulls[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
@@ -45,6 +49,8 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	MemoryContext child;
 	const char *name;
 	const char *ident;
+	int  current_context_id = (*context_id)++;
+	Size total_bytes_children = 0;
 
 	Assert(MemoryContextIsValid(context));
 
@@ -103,13 +109,28 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	values[6] = Int64GetDatum(stat.freespace);
 	values[7] = Int64GetDatum(stat.freechunks);
 	values[8] = Int64GetDatum(stat.totalspace - stat.freespace);
-	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+	values[9] = Int32GetDatum(current_context_id);
+
+	if (path == NIL)
+		nulls[10] = true;
+	else
+		values[10] = convert_path_to_datum(path);
 
+	path = lappend_int(path, current_context_id);
 	for (child = context->firstchild; child != NULL; child = child->nextchild)
 	{
-		PutMemoryContextsStatsTupleStore(tupstore, tupdesc,
-										 child, name, level + 1);
+		PutMemoryContextsStatsTupleStore(tupstore, tupdesc, child, name,
+										 level+1, context_id, path,
+										 &total_bytes_children);
 	}
+	path = list_delete_last(path);
+
+	total_bytes_children += stat.totalspace;
+	values[11] = Int64GetDatum(total_bytes_children);
+
+	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+
+	*total_bytes_inc_children += total_bytes_children;
 }
 
 /*
@@ -120,10 +141,14 @@ Datum
 pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 {
 	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	int context_id = 0;
+	List *path = NIL;
+	Size total_bytes_inc_children = 0;
 
 	InitMaterializedSRF(fcinfo, 0);
 	PutMemoryContextsStatsTupleStore(rsinfo->setResult, rsinfo->setDesc,
-									 TopMemoryContext, NULL, 0);
+									 TopMemoryContext, NULL, 0, &context_id,
+									 path, &total_bytes_inc_children);
 
 	return (Datum) 0;
 }
@@ -193,3 +218,26 @@ pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
 
 	PG_RETURN_BOOL(true);
 }
+
+/*
+ * Convert a list of context ids to a int[] Datum
+ */
+static Datum
+convert_path_to_datum(List *path)
+{
+	Datum	   *datum_array;
+	int			length;
+	ArrayType  *result_array;
+	ListCell   *lc;
+
+	length = list_length(path);
+	datum_array = (Datum *) palloc(length * sizeof(Datum));
+	length = 0;
+	foreach(lc, path)
+	{
+		datum_array[length++] = Int32GetDatum(lfirst_int(lc));
+	}
+	result_array = construct_array_builtin(datum_array, length, INT4OID);
+
+	return PointerGetDatum(result_array);
+}
\ No newline at end of file
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 5b67784731..51280c563d 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8263,9 +8263,9 @@
   proname => 'pg_get_backend_memory_contexts', prorows => '100',
   proretset => 't', provolatile => 'v', proparallel => 'r',
   prorettype => 'record', proargtypes => '',
-  proallargtypes => '{text,text,text,int4,int8,int8,int8,int8,int8}',
-  proargmodes => '{o,o,o,o,o,o,o,o,o}',
-  proargnames => '{name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes}',
+  proallargtypes => '{text,text,text,int4,int8,int8,int8,int8,int8,int4,_int4,int8}',
+  proargmodes => '{o,o,o,o,o,o,o,o,o,o,o,o}',
+  proargnames => '{name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes, id, path, total_bytes_including_children}',
   prosrc => 'pg_get_backend_memory_contexts' },
 
 # logging memory contexts of the specified backend
diff --git a/src/test/regress/expected/rules.out b/src/test/regress/expected/rules.out
index f645e8486b..cf28820b26 100644
--- a/src/test/regress/expected/rules.out
+++ b/src/test/regress/expected/rules.out
@@ -1311,8 +1311,11 @@ pg_backend_memory_contexts| SELECT name,
     total_nblocks,
     free_bytes,
     free_chunks,
-    used_bytes
-   FROM pg_get_backend_memory_contexts() pg_get_backend_memory_contexts(name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes);
+    used_bytes,
+    id,
+    path,
+    total_bytes_including_children
+   FROM pg_get_backend_memory_contexts() pg_get_backend_memory_contexts(name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes, id, path, total_bytes_including_children);
 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 271313ebf8..de9f1ce62c 100644
--- a/src/test/regress/expected/sysviews.out
+++ b/src/test/regress/expected/sysviews.out
@@ -20,7 +20,7 @@ select count(*) >= 0 as ok from pg_available_extensions;
 (1 row)
 
 -- The entire output of pg_backend_memory_contexts is not stable,
--- we test only the existence and basic condition of TopMemoryContext.
+-- we test the existence and basic condition of TopMemoryContext.
 select name, ident, parent, level, total_bytes >= free_bytes
   from pg_backend_memory_contexts where level = 0;
        name       | ident | parent | level | ?column? 
@@ -28,6 +28,27 @@ select name, ident, parent, level, total_bytes >= free_bytes
  TopMemoryContext |       |        |     0 | t
 (1 row)
 
+-- 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(*) > 0
+from contexts
+where array[(select id from contexts where name = 'CacheMemoryContext')] <@ path;
+ ?column? 
+----------
+ t
+(1 row)
+
+-- TopMemoryContext should have the largest total_bytes_including_children.
+select name from pg_backend_memory_contexts
+  order by total_bytes_including_children desc limit 1;
+       name       
+------------------
+ TopMemoryContext
+(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 6b4e24601d..1573e5012d 100644
--- a/src/test/regress/sql/sysviews.sql
+++ b/src/test/regress/sql/sysviews.sql
@@ -13,10 +13,23 @@ select count(*) >= 0 as ok from pg_available_extension_versions;
 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.
+-- we test the existence and basic condition of TopMemoryContext.
 select name, ident, parent, level, total_bytes >= free_bytes
   from pg_backend_memory_contexts where level = 0;
 
+-- 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(*) > 0
+from contexts
+where array[(select id from contexts where name = 'CacheMemoryContext')] <@ path;
+
+-- TopMemoryContext should have the largest total_bytes_including_children.
+select name from pg_backend_memory_contexts
+  order by total_bytes_including_children desc limit 1;
+
 -- At introduction, pg_config had 23 entries; it may grow
 select count(*) > 20 as ok from pg_config;
 
-- 
2.34.1

From 5a493bf26a27ed6ec9a465adedee7a0678f0138f Mon Sep 17 00:00:00 2001
From: Melih Mutlu <m.melihmutlu@gmail.com>
Date: Tue, 13 Jun 2023 15:00:05 +0300
Subject: [PATCH v5] Adding id/parent_id into pg_backend_memory_contexts

Added three new fields 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 id and parent_id are needed for each context. Those new id's 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 differenct
pg_get_backend_memory_contexts() calls.

Context id's are start from 0 which is assigned to TopMemoryContext.
---
 doc/src/sgml/system-views.sgml         | 30 ++++++++++
 src/backend/utils/adt/mcxtfuncs.c      | 81 +++++++++++++++++++++-----
 src/include/catalog/pg_proc.dat        |  6 +-
 src/test/regress/expected/rules.out    |  5 +-
 src/test/regress/expected/sysviews.out | 23 +++++++-
 src/test/regress/sql/sysviews.sql      | 15 ++++-
 6 files changed, 138 insertions(+), 22 deletions(-)

diff --git a/doc/src/sgml/system-views.sgml b/doc/src/sgml/system-views.sgml
index 72d01fc624..1d7d94b61d 100644
--- a/doc/src/sgml/system-views.sgml
+++ b/doc/src/sgml/system-views.sgml
@@ -472,6 +472,16 @@
       </para></entry>
      </row>
 
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>context_id</structfield> <type>int4</type>
+      </para>
+      <para>
+       Current context id. Note that the context id is a temporary id and may
+       change in each invocation
+      </para></entry>
+     </row>
+
      <row>
       <entry role="catalog_table_entry"><para role="column_definition">
        <structfield>ident</structfield> <type>text</type>
@@ -499,6 +509,17 @@
       </para></entry>
      </row>
 
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>path</structfield> <type>int4[]</type>
+      </para>
+      <para>
+       Path to reach the current context from TopMemoryContext. Context ids in
+       this list represents all parents of the current context. This can be
+       used to build the parent and child relation
+      </para></entry>
+     </row>
+
      <row>
       <entry role="catalog_table_entry"><para role="column_definition">
        <structfield>total_bytes</structfield> <type>int8</type>
@@ -508,6 +529,15 @@
       </para></entry>
      </row>
 
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>total_bytes_including_children</structfield> <type>int8</type>
+      </para>
+      <para>
+       Total bytes allocated for this memory context including its children
+      </para></entry>
+     </row>
+
      <row>
       <entry role="catalog_table_entry"><para role="column_definition">
        <structfield>total_nblocks</structfield> <type>int8</type>
diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
index 4708d73f5f..44e6b87fe0 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -20,6 +20,7 @@
 #include "mb/pg_wchar.h"
 #include "storage/proc.h"
 #include "storage/procarray.h"
+#include "utils/array.h"
 #include "utils/builtins.h"
 
 /* ----------
@@ -28,6 +29,8 @@
  */
 #define MEMORY_CONTEXT_IDENT_DISPLAY_SIZE	1024
 
+static Datum convert_path_to_datum(List *path);
+
 /*
  * PutMemoryContextsStatsTupleStore
  *		One recursion level for pg_get_backend_memory_contexts.
@@ -35,9 +38,10 @@
 static void
 PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 								 TupleDesc tupdesc, MemoryContext context,
-								 const char *parent, int level)
+								 const char *parent, int level, int *context_id,
+								 List *path, Size *total_bytes_inc_children)
 {
-#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	9
+#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	12
 
 	Datum		values[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
 	bool		nulls[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
@@ -45,6 +49,8 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	MemoryContext child;
 	const char *name;
 	const char *ident;
+	int  current_context_id = (*context_id)++;
+	Size total_bytes_children = 0;
 
 	Assert(MemoryContextIsValid(context));
 
@@ -73,6 +79,8 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	else
 		nulls[0] = true;
 
+	values[1] = Int32GetDatum(current_context_id);
+
 	if (ident)
 	{
 		int			idlen = strlen(ident);
@@ -87,29 +95,44 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 
 		memcpy(clipped_ident, ident, idlen);
 		clipped_ident[idlen] = '\0';
-		values[1] = CStringGetTextDatum(clipped_ident);
+		values[2] = CStringGetTextDatum(clipped_ident);
 	}
 	else
-		nulls[1] = true;
+		nulls[2] = true;
 
 	if (parent)
-		values[2] = CStringGetTextDatum(parent);
+		values[3] = CStringGetTextDatum(parent);
 	else
-		nulls[2] = true;
+		nulls[3] = true;
 
-	values[3] = Int32GetDatum(level);
-	values[4] = Int64GetDatum(stat.totalspace);
-	values[5] = Int64GetDatum(stat.nblocks);
-	values[6] = Int64GetDatum(stat.freespace);
-	values[7] = Int64GetDatum(stat.freechunks);
-	values[8] = Int64GetDatum(stat.totalspace - stat.freespace);
-	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+	values[4] = Int32GetDatum(level);
 
+	if (path == NIL)
+		nulls[5] = true;
+	else
+		values[5] = convert_path_to_datum(path);
+
+	values[6] = Int64GetDatum(stat.totalspace);
+
+	path = lappend_int(path, current_context_id);
 	for (child = context->firstchild; child != NULL; child = child->nextchild)
 	{
-		PutMemoryContextsStatsTupleStore(tupstore, tupdesc,
-										 child, name, level + 1);
+		PutMemoryContextsStatsTupleStore(tupstore, tupdesc, child, name,
+										 level+1, context_id, path,
+										 &total_bytes_children);
 	}
+	path = list_delete_last(path);
+
+	total_bytes_children += stat.totalspace;
+	values[7] = Int64GetDatum(total_bytes_children);
+	values[8] = Int64GetDatum(stat.nblocks);
+	values[9] = Int64GetDatum(stat.freespace);
+	values[10] = Int64GetDatum(stat.freechunks);
+	values[11] = Int64GetDatum(stat.totalspace - stat.freespace);
+
+	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+
+	*total_bytes_inc_children += total_bytes_children;
 }
 
 /*
@@ -120,10 +143,14 @@ Datum
 pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 {
 	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	int context_id = 0;
+	List *path = NIL;
+	Size total_bytes_inc_children = 0;
 
 	InitMaterializedSRF(fcinfo, 0);
 	PutMemoryContextsStatsTupleStore(rsinfo->setResult, rsinfo->setDesc,
-									 TopMemoryContext, NULL, 0);
+									 TopMemoryContext, NULL, 0, &context_id,
+									 path, &total_bytes_inc_children);
 
 	return (Datum) 0;
 }
@@ -193,3 +220,25 @@ pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
 
 	PG_RETURN_BOOL(true);
 }
+
+/*
+ * Convert a list of context ids to a int[] Datum
+ */
+static Datum
+convert_path_to_datum(List *path)
+{
+	Datum	   *datum_array;
+	int			length;
+	ArrayType  *result_array;
+
+	length = list_length(path);
+	datum_array = (Datum *) palloc(length * sizeof(Datum));
+	length = 0;
+	foreach_int(id, path)
+	{
+		datum_array[length++] = Int32GetDatum(id);
+	}
+	result_array = construct_array_builtin(datum_array, length, INT4OID);
+
+	return PointerGetDatum(result_array);
+}
\ No newline at end of file
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 58811a6530..3299a6cfef 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8263,9 +8263,9 @@
   proname => 'pg_get_backend_memory_contexts', prorows => '100',
   proretset => 't', provolatile => 'v', proparallel => 'r',
   prorettype => 'record', proargtypes => '',
-  proallargtypes => '{text,text,text,int4,int8,int8,int8,int8,int8}',
-  proargmodes => '{o,o,o,o,o,o,o,o,o}',
-  proargnames => '{name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes}',
+  proallargtypes => '{text,int4,text,text,int4,_int4,int8,int8,int8,int8,int8,int8}',
+  proargmodes => '{o,o,o,o,o,o,o,o,o,o,o,o}',
+  proargnames => '{name, context_id, ident, parent, level, path, total_bytes, total_bytes_including_children, 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/test/regress/expected/rules.out b/src/test/regress/expected/rules.out
index 55f2e95352..e183d1185e 100644
--- a/src/test/regress/expected/rules.out
+++ b/src/test/regress/expected/rules.out
@@ -1304,15 +1304,18 @@ pg_available_extensions| SELECT e.name,
    FROM (pg_available_extensions() e(name, default_version, comment)
      LEFT JOIN pg_extension x ON ((e.name = x.extname)));
 pg_backend_memory_contexts| SELECT name,
+    context_id,
     ident,
     parent,
     level,
+    path,
     total_bytes,
+    total_bytes_including_children,
     total_nblocks,
     free_bytes,
     free_chunks,
     used_bytes
-   FROM pg_get_backend_memory_contexts() pg_get_backend_memory_contexts(name, ident, parent, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes);
+   FROM pg_get_backend_memory_contexts() pg_get_backend_memory_contexts(name, context_id, ident, parent, level, path, total_bytes, total_bytes_including_children, 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 271313ebf8..ad45404735 100644
--- a/src/test/regress/expected/sysviews.out
+++ b/src/test/regress/expected/sysviews.out
@@ -20,7 +20,7 @@ select count(*) >= 0 as ok from pg_available_extensions;
 (1 row)
 
 -- The entire output of pg_backend_memory_contexts is not stable,
--- we test only the existence and basic condition of TopMemoryContext.
+-- we test the existence and basic condition of TopMemoryContext.
 select name, ident, parent, level, total_bytes >= free_bytes
   from pg_backend_memory_contexts where level = 0;
        name       | ident | parent | level | ?column? 
@@ -28,6 +28,27 @@ select name, ident, parent, level, total_bytes >= free_bytes
  TopMemoryContext |       |        |     0 | t
 (1 row)
 
+-- 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(*) > 0
+from contexts
+where array[(select context_id from contexts where name = 'CacheMemoryContext')] <@ path;
+ ?column? 
+----------
+ t
+(1 row)
+
+-- TopMemoryContext should have the largest total_bytes_including_children.
+select name from pg_backend_memory_contexts
+  order by total_bytes_including_children desc limit 1;
+       name       
+------------------
+ TopMemoryContext
+(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 6b4e24601d..64ecbc30cd 100644
--- a/src/test/regress/sql/sysviews.sql
+++ b/src/test/regress/sql/sysviews.sql
@@ -13,10 +13,23 @@ select count(*) >= 0 as ok from pg_available_extension_versions;
 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.
+-- we test the existence and basic condition of TopMemoryContext.
 select name, ident, parent, level, total_bytes >= free_bytes
   from pg_backend_memory_contexts where level = 0;
 
+-- 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(*) > 0
+from contexts
+where array[(select context_id from contexts where name = 'CacheMemoryContext')] <@ path;
+
+-- TopMemoryContext should have the largest total_bytes_including_children.
+select name from pg_backend_memory_contexts
+  order by total_bytes_including_children desc limit 1;
+
 -- At introduction, pg_config had 23 entries; it may grow
 select count(*) > 20 as ok from pg_config;
 
-- 
2.34.1

From 39e4bb039900c270b42475648c8bb0ae64a3f14c Mon Sep 17 00:00:00 2001
From: Melih Mutlu <m.melihmutlu@gmail.com>
Date: Mon, 1 Jul 2024 16:06:03 +0300
Subject: [PATCH v6 2/2] Add total_bytes_including_children column

Add a new column total_bytes_including_children into
pg_backend_memory_contexts.
---
 doc/src/sgml/system-views.sgml         |  9 +++++++++
 src/backend/utils/adt/mcxtfuncs.c      | 28 +++++++++++++++++---------
 src/include/catalog/pg_proc.dat        |  6 +++---
 src/test/regress/expected/rules.out    |  3 ++-
 src/test/regress/expected/sysviews.out | 14 +++++++++++++
 src/test/regress/sql/sysviews.sql      | 10 +++++++++
 6 files changed, 56 insertions(+), 14 deletions(-)

diff --git a/doc/src/sgml/system-views.sgml b/doc/src/sgml/system-views.sgml
index f818aba974..ad2cb71a88 100644
--- a/doc/src/sgml/system-views.sgml
+++ b/doc/src/sgml/system-views.sgml
@@ -530,6 +530,15 @@
       </para></entry>
      </row>
 
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>total_bytes_including_children</structfield> <type>int8</type>
+      </para>
+      <para>
+       Total bytes allocated for this memory context including its children
+      </para></entry>
+     </row>
+
      <row>
       <entry role="catalog_table_entry"><para role="column_definition">
        <structfield>total_nblocks</structfield> <type>int8</type>
diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
index 82b863341c..ca1eb0840e 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -38,9 +38,9 @@ static void
 PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 								 TupleDesc tupdesc, MemoryContext context,
 								 const char *parent, int level, int *context_id,
-								 List *path)
+								 List *path, Size *total_bytes_inc_children)
 {
-#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	11
+#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	12
 
 	Datum		values[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
 	bool		nulls[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
@@ -50,6 +50,7 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	const char *ident;
 	const char *type;
 	int  current_context_id = (*context_id)++;
+	Size total_bytes_children = 0;
 
 	Assert(MemoryContextIsValid(context));
 
@@ -127,20 +128,26 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	path = lappend_int(path, current_context_id);
 	values[5] = convert_path_to_datum(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);
-
+	total_bytes_children += stat.totalspace;
 	for (child = context->firstchild; child != NULL; child = child->nextchild)
 	{
 		PutMemoryContextsStatsTupleStore(tupstore, tupdesc, child, name,
-										 level+1, context_id, path);
+										 level+1, context_id, path,
+										 &total_bytes_children);
 	}
 	path = list_delete_last(path);
+
+	values[6] = Int64GetDatum(stat.totalspace);
+	values[7] = Int64GetDatum(total_bytes_children);
+	values[8] = Int64GetDatum(stat.nblocks);
+	values[9] = Int64GetDatum(stat.freespace);
+	values[10] = Int64GetDatum(stat.freechunks);
+	values[11] = Int64GetDatum(stat.totalspace - stat.freespace);
+
 	
 	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+
+	*total_bytes_inc_children += total_bytes_children;
 }
 
 /*
@@ -153,11 +160,12 @@ pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
 	int context_id = 0;
 	List *path = NIL;
+	Size total_bytes_inc_children = 0;
 
 	InitMaterializedSRF(fcinfo, 0);
 	PutMemoryContextsStatsTupleStore(rsinfo->setResult, rsinfo->setDesc,
 									 TopMemoryContext, NULL, 0, &context_id,
-									 path);
+									 path, &total_bytes_inc_children);
 
 	return (Datum) 0;
 }
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 202a0b7567..be31190f62 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8279,9 +8279,9 @@
   proname => 'pg_get_backend_memory_contexts', prorows => '100',
   proretset => 't', provolatile => 'v', proparallel => 'r',
   prorettype => 'record', proargtypes => '',
-  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}',
+  proallargtypes => '{text,text,text,text,int4,_int4,int8,int8,int8,int8,int8,int8}',
+  proargmodes => '{o,o,o,o,o,o,o,o,o,o,o,o}',
+  proargnames => '{name, ident, parent, type, level, path, total_bytes, total_bytes_including_children, 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/test/regress/expected/rules.out b/src/test/regress/expected/rules.out
index 5201280669..9318cb5212 100644
--- a/src/test/regress/expected/rules.out
+++ b/src/test/regress/expected/rules.out
@@ -1310,11 +1310,12 @@ pg_backend_memory_contexts| SELECT name,
     level,
     path,
     total_bytes,
+    total_bytes_including_children,
     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);
+   FROM pg_get_backend_memory_contexts() pg_get_backend_memory_contexts(name, ident, parent, type, level, path, total_bytes, total_bytes_including_children, 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 b27b7dca4b..fa64b6de9b 100644
--- a/src/test/regress/expected/sysviews.out
+++ b/src/test/regress/expected/sysviews.out
@@ -64,6 +64,20 @@ where array[(select path[level+1] from contexts where name = 'CacheMemoryContext
  t
 (1 row)
 
+-- Test whether total_bytes_including_children really sums up to the total
+-- bytes used by all child contexts of CacheMemoryContext.
+with contexts as (
+  select * from pg_backend_memory_contexts
+)
+select sum(total_bytes) = 
+      (select total_bytes_including_children from contexts where name = 'CacheMemoryContext')
+from contexts
+where array[(select path[level+1] from contexts where name = 'CacheMemoryContext')] <@ path;
+ ?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 0953d3e2c7..91da7eb5bb 100644
--- a/src/test/regress/sql/sysviews.sql
+++ b/src/test/regress/sql/sysviews.sql
@@ -41,6 +41,16 @@ select count(*) > 0
 from contexts
 where array[(select path[level+1] from contexts where name = 'CacheMemoryContext')] <@ path;
 
+-- Test whether total_bytes_including_children really sums up to the total
+-- bytes used by all child contexts of CacheMemoryContext.
+with contexts as (
+  select * from pg_backend_memory_contexts
+)
+select sum(total_bytes) = 
+      (select total_bytes_including_children from contexts where name = 'CacheMemoryContext')
+from contexts
+where array[(select path[level+1] from contexts where name = 'CacheMemoryContext')] <@ path;
+
 -- At introduction, pg_config had 23 entries; it may grow
 select count(*) > 20 as ok from pg_config;
 
-- 
2.34.1

From 54086bd9923322459e1484a45b30c9bf934ae7ac Mon Sep 17 00:00:00 2001
From: Melih Mutlu <m.melihmutlu@gmail.com>
Date: Mon, 1 Jul 2024 13:00:51 +0300
Subject: [PATCH v6 1/2] 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         | 32 ++++++++++++++
 src/backend/utils/adt/mcxtfuncs.c      | 58 +++++++++++++++++++++-----
 src/include/catalog/pg_proc.dat        |  6 +--
 src/test/regress/expected/rules.out    |  3 +-
 src/test/regress/expected/sysviews.out | 13 ++++++
 src/test/regress/sql/sysviews.sql      |  9 ++++
 6 files changed, 106 insertions(+), 15 deletions(-)

diff --git a/doc/src/sgml/system-views.sgml b/doc/src/sgml/system-views.sgml
index bdc34cf94e..f818aba974 100644
--- a/doc/src/sgml/system-views.sgml
+++ b/doc/src/sgml/system-views.sgml
@@ -508,6 +508,19 @@
       </para></entry>
      </row>
 
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>path</structfield> <type>int4[]</type>
+      </para>
+      <para>
+       Path that includes all contexts from TopMemoryContext to the current
+       context. Context IDs in this list represents all parents of a context,
+       and last element in the list refers to the context itself. This can be
+       used to build the parent and child relation. Note that IDs used in path
+       are transient and may change in each execution
+      </para></entry>
+     </row>
+
      <row>
       <entry role="catalog_table_entry"><para role="column_definition">
        <structfield>total_bytes</structfield> <type>int8</type>
@@ -561,6 +574,25 @@
    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 a memory context and its
+   child contexts:
+<programlisting>
+WITH memory_contexts AS (
+    SELECT * 
+    FROM pg_backend_memory_contexts
+)
+SELECT SUM(total_bytes)
+FROM memory_contexts
+WHERE ARRAY[(SELECT path[array_length(path, 1)] FROM memory_contexts WHERE name = 'CacheMemoryContext')] &lt;@ path;
+</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..82b863341c 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -19,6 +19,7 @@
 #include "mb/pg_wchar.h"
 #include "storage/proc.h"
 #include "storage/procarray.h"
+#include "utils/array.h"
 #include "utils/builtins.h"
 
 /* ----------
@@ -27,6 +28,8 @@
  */
 #define MEMORY_CONTEXT_IDENT_DISPLAY_SIZE	1024
 
+static Datum convert_path_to_datum(List *path);
+
 /*
  * PutMemoryContextsStatsTupleStore
  *		One recursion level for pg_get_backend_memory_contexts.
@@ -34,9 +37,10 @@
 static void
 PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 								 TupleDesc tupdesc, MemoryContext context,
-								 const char *parent, int level)
+								 const char *parent, int level, int *context_id,
+								 List *path)
 {
-#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];
@@ -45,6 +49,7 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	const char *name;
 	const char *ident;
 	const char *type;
+	int  current_context_id = (*context_id)++;
 
 	Assert(MemoryContextIsValid(context));
 
@@ -118,18 +123,24 @@ 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);
+
+	path = lappend_int(path, current_context_id);
+	values[5] = convert_path_to_datum(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);
+		PutMemoryContextsStatsTupleStore(tupstore, tupdesc, child, name,
+										 level+1, context_id, path);
 	}
+	path = list_delete_last(path);
+	
+	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
 }
 
 /*
@@ -140,10 +151,13 @@ Datum
 pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 {
 	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	int context_id = 0;
+	List *path = NIL;
 
 	InitMaterializedSRF(fcinfo, 0);
 	PutMemoryContextsStatsTupleStore(rsinfo->setResult, rsinfo->setDesc,
-									 TopMemoryContext, NULL, 0);
+									 TopMemoryContext, NULL, 0, &context_id,
+									 path);
 
 	return (Datum) 0;
 }
@@ -206,3 +220,25 @@ pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
 
 	PG_RETURN_BOOL(true);
 }
+
+/*
+ * Convert a list of context ids to a int[] Datum
+ */
+static Datum
+convert_path_to_datum(List *path)
+{
+	Datum	   *datum_array;
+	int			length;
+	ArrayType  *result_array;
+
+	length = list_length(path);
+	datum_array = (Datum *) palloc(length * sizeof(Datum));
+	length = 0;
+	foreach_int(id, path)
+	{
+		datum_array[length++] = Int32GetDatum(id);
+	}
+	result_array = construct_array_builtin(datum_array, length, INT4OID);
+
+	return PointerGetDatum(result_array);
+}
\ No newline at end of file
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index d4ac578ae6..202a0b7567 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8279,9 +8279,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/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..b27b7dca4b 100644
--- a/src/test/regress/expected/sysviews.out
+++ b/src/test/regress/expected/sysviews.out
@@ -51,6 +51,19 @@ 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(*) > 0
+from contexts
+where array[(select path[level+1] from contexts where name = 'CacheMemoryContext')] <@ path;
+ ?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..0953d3e2c7 100644
--- a/src/test/regress/sql/sysviews.sql
+++ b/src/test/regress/sql/sysviews.sql
@@ -32,6 +32,15 @@ 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(*) > 0
+from contexts
+where array[(select path[level+1] from contexts where name = 'CacheMemoryContext')] <@ path;
+
 -- At introduction, pg_config had 23 entries; it may grow
 select count(*) > 20 as ok from pg_config;
 
-- 
2.34.1

From e448d8b3a941ecab539bdc57151be9f3968ca3de Mon Sep 17 00:00:00 2001
From: Melih Mutlu <m.melihmutlu@gmail.com>
Date: Mon, 1 Jul 2024 16:06:03 +0300
Subject: [PATCH v7 2/2] Add total_bytes_including_children column

Add a new column total_bytes_including_children into
pg_backend_memory_contexts.
---
 doc/src/sgml/system-views.sgml         |  9 +++++++++
 src/backend/utils/adt/mcxtfuncs.c      | 25 ++++++++++++++++---------
 src/include/catalog/pg_proc.dat        |  6 +++---
 src/test/regress/expected/rules.out    |  3 ++-
 src/test/regress/expected/sysviews.out | 14 ++++++++++++++
 src/test/regress/sql/sysviews.sql      | 10 ++++++++++
 6 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/doc/src/sgml/system-views.sgml b/doc/src/sgml/system-views.sgml
index 501244709d..0d6ad8e55a 100644
--- a/doc/src/sgml/system-views.sgml
+++ b/doc/src/sgml/system-views.sgml
@@ -531,6 +531,15 @@
       </para></entry>
      </row>
 
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>total_bytes_including_children</structfield> <type>int8</type>
+      </para>
+      <para>
+       Total bytes allocated for this memory context including its children
+      </para></entry>
+     </row>
+
      <row>
       <entry role="catalog_table_entry"><para role="column_definition">
        <structfield>total_nblocks</structfield> <type>int8</type>
diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
index 98f03b3279..db947cdc3e 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -38,9 +38,10 @@ static void
 PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 								 TupleDesc tupdesc, MemoryContext context,
 								 const char *parent, int level, int *context_id,
-								 List *context_ids)
+								 List *context_ids,
+								 Size *total_bytes_inc_children)
 {
-#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	11
+#define PG_GET_BACKEND_MEMORY_CONTEXTS_COLS	12
 
 	Datum		values[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
 	bool		nulls[PG_GET_BACKEND_MEMORY_CONTEXTS_COLS];
@@ -50,6 +51,7 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	const char *ident;
 	const char *type;
 	int  current_context_id = (*context_id)++;
+	Size total_bytes_children = 0;
 
 	Assert(MemoryContextIsValid(context));
 
@@ -127,21 +129,25 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	context_ids = lappend_int(context_ids, current_context_id);
 	values[5] = convert_ids_to_datum(context_ids);
 
+	total_bytes_children += stat.totalspace;
 	for (child = context->firstchild; child != NULL; child = child->nextchild)
 	{
 		PutMemoryContextsStatsTupleStore(tupstore, tupdesc, child, name,
-										 level+1, context_id, context_ids);
+										 level+1, context_id, context_ids,
+										 &total_bytes_children);
 	}
 	context_ids = list_delete_last(context_ids);
 
 	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);
-
+	values[7] = Int64GetDatum(total_bytes_children);
+	values[8] = Int64GetDatum(stat.nblocks);
+	values[9] = Int64GetDatum(stat.freespace);
+	values[10] = Int64GetDatum(stat.freechunks);
+	values[11] = Int64GetDatum(stat.totalspace - stat.freespace);
 
 	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+
+	*total_bytes_inc_children += total_bytes_children;
 }
 
 /*
@@ -154,11 +160,12 @@ pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
 	int context_id = 0;
 	List *context_ids = NIL;
+	Size total_bytes_inc_children = 0;
 
 	InitMaterializedSRF(fcinfo, 0);
 	PutMemoryContextsStatsTupleStore(rsinfo->setResult, rsinfo->setDesc,
 									 TopMemoryContext, NULL, 0, &context_id,
-									 context_ids);
+									 context_ids, &total_bytes_inc_children);
 
 	return (Datum) 0;
 }
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 18d0187a80..8bf8253ab7 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8284,9 +8284,9 @@
   proname => 'pg_get_backend_memory_contexts', prorows => '100',
   proretset => 't', provolatile => 'v', proparallel => 'r',
   prorettype => 'record', proargtypes => '',
-  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, context_ids, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes}',
+  proallargtypes => '{text,text,text,text,int4,_int4,int8,int8,int8,int8,int8,int8}',
+  proargmodes => '{o,o,o,o,o,o,o,o,o,o,o,o}',
+  proargnames => '{name, ident, parent, type, level, context_ids, total_bytes, total_bytes_including_children, 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/test/regress/expected/rules.out b/src/test/regress/expected/rules.out
index 7593e94939..35ffd387c7 100644
--- a/src/test/regress/expected/rules.out
+++ b/src/test/regress/expected/rules.out
@@ -1310,11 +1310,12 @@ pg_backend_memory_contexts| SELECT name,
     level,
     context_ids,
     total_bytes,
+    total_bytes_including_children,
     total_nblocks,
     free_bytes,
     free_chunks,
     used_bytes
-   FROM pg_get_backend_memory_contexts() pg_get_backend_memory_contexts(name, ident, parent, type, level, context_ids, 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, context_ids, total_bytes, total_bytes_including_children, 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 af4d82adf4..4bec9a752d 100644
--- a/src/test/regress/expected/sysviews.out
+++ b/src/test/regress/expected/sysviews.out
@@ -65,6 +65,20 @@ where c1.name = 'CacheMemoryContext'
  t
 (1 row)
 
+-- Test whether total_bytes_including_children really sums up to the total
+-- bytes used by all child contexts of CacheMemoryContext.
+with contexts as (
+  select * from pg_backend_memory_contexts
+)
+select min(c1.total_bytes_including_children) = sum(c2.total_bytes)
+from contexts c1, contexts c2
+where c1.name = 'CacheMemoryContext'
+  and c1.context_ids[1] = c2.context_ids[c2.level-c1.level+1];
+ ?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 f0ac59387d..ae5fd1a9fa 100644
--- a/src/test/regress/sql/sysviews.sql
+++ b/src/test/regress/sql/sysviews.sql
@@ -42,6 +42,16 @@ from contexts c1, contexts c2
 where c1.name = 'CacheMemoryContext'
   and c1.context_ids[1] = c2.context_ids[c2.level-c1.level+1];
 
+-- Test whether total_bytes_including_children really sums up to the total
+-- bytes used by all child contexts of CacheMemoryContext.
+with contexts as (
+  select * from pg_backend_memory_contexts
+)
+select min(c1.total_bytes_including_children) = sum(c2.total_bytes)
+from contexts c1, contexts c2
+where c1.name = 'CacheMemoryContext'
+  and c1.context_ids[1] = c2.context_ids[c2.level-c1.level+1];
+
 -- At introduction, pg_config had 23 entries; it may grow
 select count(*) > 20 as ok from pg_config;
 
-- 
2.34.1

From f665b560103f322bbd74430f1dfbb458ab578c5f Mon Sep 17 00:00:00 2001
From: Melih Mutlu <m.melihmutlu@gmail.com>
Date: Mon, 1 Jul 2024 13:00:51 +0300
Subject: [PATCH v7 1/2] context_ids column in 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         | 33 ++++++++++++++
 src/backend/utils/adt/mcxtfuncs.c      | 60 +++++++++++++++++++++-----
 src/include/catalog/pg_proc.dat        |  6 +--
 src/test/regress/expected/rules.out    |  3 +-
 src/test/regress/expected/sysviews.out | 14 ++++++
 src/test/regress/sql/sysviews.sql      | 10 +++++
 6 files changed, 111 insertions(+), 15 deletions(-)

diff --git a/doc/src/sgml/system-views.sgml b/doc/src/sgml/system-views.sgml
index bdc34cf94e..501244709d 100644
--- a/doc/src/sgml/system-views.sgml
+++ b/doc/src/sgml/system-views.sgml
@@ -508,6 +508,20 @@
       </para></entry>
      </row>
 
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>context_ids</structfield> <type>int4[]</type>
+      </para>
+      <para>
+       Array of transient identifiers to describe the memory context
+       hierarchy. The first array element contains the ID for the current
+       context and each subsequent ID is the parent of the previous element.
+       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>
+
      <row>
       <entry role="catalog_table_entry"><para role="column_definition">
        <structfield>total_bytes</structfield> <type>int8</type>
@@ -561,6 +575,25 @@
    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 a memory context and its
+   child contexts:
+<programlisting>
+WITH memory_contexts AS (
+    SELECT *
+    FROM pg_backend_memory_contexts
+)
+SELECT SUM(total_bytes)
+FROM memory_contexts
+WHERE ARRAY[(SELECT context_ids[1] FROM memory_contexts WHERE name = 'CacheMemoryContext')] &lt;@ context_ids;
+</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..98f03b3279 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -19,6 +19,7 @@
 #include "mb/pg_wchar.h"
 #include "storage/proc.h"
 #include "storage/procarray.h"
+#include "utils/array.h"
 #include "utils/builtins.h"
 
 /* ----------
@@ -27,6 +28,8 @@
  */
 #define MEMORY_CONTEXT_IDENT_DISPLAY_SIZE	1024
 
+static Datum convert_ids_to_datum(List *path);
+
 /*
  * PutMemoryContextsStatsTupleStore
  *		One recursion level for pg_get_backend_memory_contexts.
@@ -34,9 +37,10 @@
 static void
 PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 								 TupleDesc tupdesc, MemoryContext context,
-								 const char *parent, int level)
+								 const char *parent, int level, int *context_id,
+								 List *context_ids)
 {
-#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];
@@ -45,6 +49,7 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	const char *name;
 	const char *ident;
 	const char *type;
+	int  current_context_id = (*context_id)++;
 
 	Assert(MemoryContextIsValid(context));
 
@@ -118,18 +123,25 @@ 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);
+
+	context_ids = lappend_int(context_ids, current_context_id);
+	values[5] = convert_ids_to_datum(context_ids);
 
 	for (child = context->firstchild; child != NULL; child = child->nextchild)
 	{
-		PutMemoryContextsStatsTupleStore(tupstore, tupdesc,
-										 child, name, level + 1);
+		PutMemoryContextsStatsTupleStore(tupstore, tupdesc, child, name,
+										 level+1, context_id, context_ids);
 	}
+	context_ids = list_delete_last(context_ids);
+
+	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);
+
+
+	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
 }
 
 /*
@@ -140,10 +152,13 @@ Datum
 pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 {
 	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	int context_id = 0;
+	List *context_ids = NIL;
 
 	InitMaterializedSRF(fcinfo, 0);
 	PutMemoryContextsStatsTupleStore(rsinfo->setResult, rsinfo->setDesc,
-									 TopMemoryContext, NULL, 0);
+									 TopMemoryContext, NULL, 0, &context_id,
+									 context_ids);
 
 	return (Datum) 0;
 }
@@ -206,3 +221,26 @@ pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
 
 	PG_RETURN_BOOL(true);
 }
+
+/*
+ * Convert a list of context ids to a int[] Datum
+ */
+static Datum
+convert_ids_to_datum(List *context_ids)
+{
+	Datum	   *datum_array;
+	int			length;
+	ArrayType  *result_array;
+	int			pos;
+
+	length = list_length(context_ids);
+	datum_array = (Datum *) palloc(length * sizeof(Datum));
+	pos = length;
+	foreach_int(id, context_ids)
+	{
+		datum_array[--pos] = Int32GetDatum(id);
+	}
+	result_array = construct_array_builtin(datum_array, length, INT4OID);
+
+	return PointerGetDatum(result_array);
+}
\ No newline at end of file
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index e899ed5e77..18d0187a80 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8284,9 +8284,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, context_ids, 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/test/regress/expected/rules.out b/src/test/regress/expected/rules.out
index 4c789279e5..7593e94939 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,
+    context_ids,
     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, context_ids, 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..af4d82adf4 100644
--- a/src/test/regress/expected/sysviews.out
+++ b/src/test/regress/expected/sysviews.out
@@ -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 c1.name = 'CacheMemoryContext'
+  and c1.context_ids[1] = c2.context_ids[c2.level-c1.level+1];
+ ?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..f0ac59387d 100644
--- a/src/test/regress/sql/sysviews.sql
+++ b/src/test/regress/sql/sysviews.sql
@@ -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 c1.name = 'CacheMemoryContext'
+  and c1.context_ids[1] = c2.context_ids[c2.level-c1.level+1];
+
 -- At introduction, pg_config had 23 entries; it may grow
 select count(*) > 20 as ok from pg_config;
 
-- 
2.34.1

diff --git a/doc/src/sgml/system-views.sgml b/doc/src/sgml/system-views.sgml
index bdc34cf94e..f818aba974 100644
--- a/doc/src/sgml/system-views.sgml
+++ b/doc/src/sgml/system-views.sgml
@@ -508,6 +508,19 @@
       </para></entry>
      </row>
 
+     <row>
+      <entry role="catalog_table_entry"><para role="column_definition">
+       <structfield>path</structfield> <type>int4[]</type>
+      </para>
+      <para>
+       Path that includes all contexts from TopMemoryContext to the current
+       context. Context IDs in this list represents all parents of a context,
+       and last element in the list refers to the context itself. This can be
+       used to build the parent and child relation. Note that IDs used in path
+       are transient and may change in each execution
+      </para></entry>
+     </row>
+
      <row>
       <entry role="catalog_table_entry"><para role="column_definition">
        <structfield>total_bytes</structfield> <type>int8</type>
@@ -561,6 +574,25 @@
    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 a memory context and its
+   child contexts:
+<programlisting>
+WITH memory_contexts AS (
+    SELECT * 
+    FROM pg_backend_memory_contexts
+)
+SELECT SUM(total_bytes)
+FROM memory_contexts
+WHERE ARRAY[(SELECT path[array_length(path, 1)] FROM memory_contexts WHERE name = 'CacheMemoryContext')] &lt;@ path;
+</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..b0b4b1aebf 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,101 @@
  */
 #define MEMORY_CONTEXT_IDENT_DISPLAY_SIZE	1024
 
+typedef struct MemoryContextId
+{
+	MemoryContext context;
+	int context_id;
+} MemoryContextId;
+
+/*
+ * get_memory_context_name_and_indent
+ *		Populate *name and *ident from the name and ident from 'context'.
+ */
+static void
+get_memory_context_name_and_indent(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_indent(context, &name, &ident);
+
 	if (name)
 		values[0] = CStringGetTextDatum(name);
 	else
@@ -75,15 +132,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 +151,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_indent(context->parent,
+										   &parent_name,
+										   &parent_ident);
+		values[2] = CStringGetTextDatum(parent_name);
+	}
 	else
 		nulls[2] = true;
 
@@ -117,19 +183,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 +203,76 @@ Datum
 pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 {
 	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	int context_id;
+	List *queue;
+	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;
+
+	while (queue != NIL)
+	{
+		List *nextQueue = NIL;
+		ListCell *lc;
+
+		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 of the outer loop.
+			 */
+			for (MemoryContext c = cur->firstchild; c != NULL; c = c->nextchild)
+				nextQueue = lappend(nextQueue, c);
+		}
+
+		/* process the next level down */
+		list_free(queue);
+		queue = nextQueue;
+	}
+
+	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 e899ed5e77..4a48f19da7 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8284,9 +8284,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..b27b7dca4b 100644
--- a/src/test/regress/expected/sysviews.out
+++ b/src/test/regress/expected/sysviews.out
@@ -51,6 +51,19 @@ 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(*) > 0
+from contexts
+where array[(select path[level+1] from contexts where name = 'CacheMemoryContext')] <@ path;
+ ?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..0953d3e2c7 100644
--- a/src/test/regress/sql/sysviews.sql
+++ b/src/test/regress/sql/sysviews.sql
@@ -32,6 +32,15 @@ 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(*) > 0
+from contexts
+where array[(select path[level+1] from contexts where name = 'CacheMemoryContext')] <@ path;
+
 -- At introduction, pg_config had 23 entries; it may grow
 select count(*) > 20 as ok from pg_config;
 

From 57376854684ab110a1aa99346dd7d53c61d3f387 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 v9] 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         |  39 ++++-
 src/backend/utils/adt/mcxtfuncs.c      | 193 ++++++++++++++++++++-----
 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 |  17 ++-
 src/test/regress/sql/sysviews.sql      |  11 +-
 7 files changed, 227 insertions(+), 46 deletions(-)

diff --git a/doc/src/sgml/system-views.sgml b/doc/src/sgml/system-views.sgml
index bdc34cf94e..b36ab4eb51 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,25 @@
    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 a memory context and its
+   child contexts:
+<programlisting>
+WITH memory_contexts AS (
+    SELECT *
+    FROM pg_backend_memory_contexts
+)
+SELECT SUM(total_bytes)
+FROM memory_contexts
+WHERE ARRAY[(SELECT path[array_length(path, 1)] FROM memory_contexts WHERE name = 'CacheMemoryContext')] &lt;@ path;
+</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..3933977b76 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,101 @@
  */
 #define MEMORY_CONTEXT_IDENT_DISPLAY_SIZE	1024
 
+typedef struct MemoryContextId
+{
+	MemoryContext context;
+	int context_id;
+} MemoryContextId;
+
+/*
+ * get_memory_context_name_and_indent
+ *		Populate *name and *ident from the name and ident from 'context'.
+ */
+static void
+get_memory_context_name_and_indent(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_indent(context, &name, &ident);
+
 	if (name)
 		values[0] = CStringGetTextDatum(name);
 	else
@@ -75,15 +132,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 +151,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_indent(context->parent,
+										   &parent_name,
+										   &parent_ident);
+		values[2] = CStringGetTextDatum(parent_name);
+	}
 	else
 		nulls[2] = true;
 
@@ -117,19 +183,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 +203,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 of the outer loop.
+		 */
+		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..5b258869d7 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,19 @@ 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(*) > 0
+from contexts
+where array[(select path[level] from contexts where name = 'CacheMemoryContext')] <@ path;
+ ?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..9874828bb6 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,15 @@ 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(*) > 0
+from contexts
+where array[(select path[level] from contexts where name = 'CacheMemoryContext')] <@ path;
+
 -- At introduction, pg_config had 23 entries; it may grow
 select count(*) > 20 as ok from pg_config;
 
-- 
2.34.1

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