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

