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

