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

