From 257b7d87aaa55c415e205afad9ba6ae90e8d6195 Mon Sep 17 00:00:00 2001
From: Daniel Gustafsson <daniel@yesql.se>
Date: Mon, 8 Dec 2025 10:38:43 +0100
Subject: [PATCH v43 3/3] Review comments

---
 doc/src/sgml/func/func-admin.sgml             |   4 +-
 src/backend/catalog/system_functions.sql      |  14 +
 src/backend/catalog/system_views.sql          |   5 -
 src/backend/utils/adt/mcxtfuncs.c             | 314 +++++++++---------
 src/include/utils/memutils.h                  |   2 +-
 .../t/001_memcontext_inj.pl                   |  42 ++-
 src/tools/pgindent/typedefs.list              |   2 +-
 7 files changed, 195 insertions(+), 188 deletions(-)

diff --git a/doc/src/sgml/func/func-admin.sgml b/doc/src/sgml/func/func-admin.sgml
index a5c66837241..3e71ced60a2 100644
--- a/doc/src/sgml/func/func-admin.sgml
+++ b/doc/src/sgml/func/func-admin.sgml
@@ -257,7 +257,7 @@
         <indexterm>
          <primary>pg_get_process_memory_contexts</primary>
         </indexterm>
-        <function>pg_get_process_memory_contexts</function> ( <parameter>pid</parameter> <type>integer</type>, <parameter>summary</parameter> <type>boolean</type> )
+        <function>pg_get_process_memory_contexts</function> ( <parameter>pid</parameter> <type>integer</type> <optional>,<parameter>summary</parameter> <type>boolean</type> <literal>DEFAULT</literal> <literal>false</literal></optional> )
         <returnvalue>setof record</returnvalue>
         ( <parameter>name</parameter> <type>text</type>,
         <parameter>ident</parameter> <type>text</type>,
@@ -360,7 +360,7 @@
         Statistics for contexts on level 2 and below are aggregates of all
         child contexts' statistics, where <literal>num_agg_contexts</literal>
         indicate the number aggregated child contexts.  When
-        <parameter>summary</parameter> is <literal>false</literal>,
+        <parameter>summary</parameter> is <literal>false</literal> (the default),
         <literal>the num_agg_contexts</literal> value is <literal>1</literal>,
         indicating that individual statistics are being displayed.
        </para>
diff --git a/src/backend/catalog/system_functions.sql b/src/backend/catalog/system_functions.sql
index 2d946d6d9e9..7b40bac5f57 100644
--- a/src/backend/catalog/system_functions.sql
+++ b/src/backend/catalog/system_functions.sql
@@ -657,6 +657,17 @@ LANGUAGE INTERNAL
 STRICT VOLATILE PARALLEL UNSAFE
 AS 'pg_replication_origin_session_setup';
 
+CREATE OR REPLACE FUNCTION
+  pg_get_process_memory_contexts(IN pid integer, IN summary boolean DEFAULT false,
+  OUT name text, OUT ident text, OUT type text, OUT level integer,
+  OUT path integer[], OUT total_bytes bigint, OUT total_nblocks bigint,
+  OUT free_bytes bigint, OUT free_chunks bigint, OUT used_bytes bigint,
+  OUT num_agg_contexts integer)
+RETURNS SETOF RECORD
+LANGUAGE INTERNAL
+STRICT VOLATILE PARALLEL UNSAFE
+AS 'pg_get_process_memory_contexts';
+
 --
 -- The default permissions for functions mean that anyone can execute them.
 -- A number of functions shouldn't be executable by just anyone, but rather
@@ -782,6 +793,7 @@ REVOKE EXECUTE ON FUNCTION pg_ls_logicalmapdir() FROM PUBLIC;
 
 REVOKE EXECUTE ON FUNCTION pg_ls_replslotdir(text) FROM PUBLIC;
 
+REVOKE EXECUTE ON FUNCTION pg_get_process_memory_contexts(integer, boolean) FROM PUBLIC;
 --
 -- We also set up some things as accessible to standard roles.
 --
@@ -808,6 +820,8 @@ GRANT EXECUTE ON FUNCTION pg_current_logfile() TO pg_monitor;
 
 GRANT EXECUTE ON FUNCTION pg_current_logfile(text) TO pg_monitor;
 
+GRANT EXECUTE ON FUNCTION pg_get_process_memory_contexts(integer, boolean) TO pg_read_all_stats;
+
 GRANT pg_read_all_settings TO pg_monitor;
 
 GRANT pg_read_all_stats TO pg_monitor;
diff --git a/src/backend/catalog/system_views.sql b/src/backend/catalog/system_views.sql
index 17ec512622b..086c4c8fb6f 100644
--- a/src/backend/catalog/system_views.sql
+++ b/src/backend/catalog/system_views.sql
@@ -692,11 +692,6 @@ GRANT SELECT ON pg_backend_memory_contexts TO pg_read_all_stats;
 REVOKE EXECUTE ON FUNCTION pg_get_backend_memory_contexts() FROM PUBLIC;
 GRANT EXECUTE ON FUNCTION pg_get_backend_memory_contexts() TO pg_read_all_stats;
 
-REVOKE EXECUTE ON FUNCTION
-	pg_get_process_memory_contexts(integer, boolean) FROM PUBLIC;
-GRANT EXECUTE ON FUNCTION
-	pg_get_process_memory_contexts(integer, boolean) TO pg_read_all_stats;
-
 -- Statistics views
 
 CREATE VIEW pg_stat_all_tables AS
diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
index c661eef7ae9..f1b5c3a0887 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -49,8 +49,11 @@
  */
 #define MAX_MEMORY_CONTEXT_STATS_NUM MEMORY_CONTEXT_REPORT_MAX_PER_BACKEND / (sizeof(MemoryStatsEntry))
 
-/* Size of dshash key */
-#define CLIENT_KEY_SIZE 32
+/*
+ * Size of dshash key. The key is a uint32 rendered as a string, 10 chars
+ * plus space for a NULL terminator can hold all the values.
+ */
+#define CLIENT_KEY_SIZE (10 + 1)
 
 /* Dynamic shared memory state for reporting statistics per context */
 typedef struct MemoryStatsEntry
@@ -92,8 +95,7 @@ static const dshash_parameters memctx_dsh_params = {
 };
 
 /*
- * These are used for reporting memory context
- * statistics of a process.
+ * These are used for reporting memory context statistics of a process.
  */
 
 /* Lock to control access to client_keys array */
@@ -103,10 +105,9 @@ static LWLock *client_keys_lock = NULL;
 static int *client_keys = NULL;
 
 /*
- * Table to store pointers to dsa memory containing
- * memory statistics and other meta data. There is one
- * entry per client backend request, keyed by ProcNumber of
- * the client obtained from client_keys array above.
+ * Table to store pointers to DSA memory containing memory statistics and other
+ * metadata. There is one entry per client backend request, keyed by ProcNumber
+ * of the client obtained from client_keys array above.
  */
 static dshash_table *MemoryStatsDsHash = NULL;
 
@@ -125,8 +126,6 @@ static void PublishMemoryContext(MemoryStatsEntry *memcxt_info,
 								 MemoryContextCounters stat,
 								 int num_contexts);
 static List *compute_context_path(MemoryContext c, HTAB *context_id_lookup);
-static void end_memorycontext_reporting(MemoryStatsDSHashEntry *entry, MemoryContext oldcontext,
-										HTAB *context_id_lookup);
 
 /* ----------
  * The max bytes for showing identifiers of MemoryContext.
@@ -140,15 +139,14 @@ static void end_memorycontext_reporting(MemoryStatsDSHashEntry *entry, MemoryCon
 #define MEMORY_STATS_MAX_TIMEOUT 5
 
 /*
- * MemoryStatsContextId
- *		Used for storage of transient identifiers for
- *		pg_get_backend_memory_contexts and the likes.
+ * MemoryContextId
+ *		Used for storage of transient identifiers for memory context reporting
  */
-typedef struct MemoryStatsContextId
+typedef struct MemoryContextId
 {
 	MemoryContext context;
 	int			context_id;
-} MemoryStatsContextId;
+} MemoryContextId;
 
 /*
  * int_list_to_array
@@ -199,7 +197,7 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	 */
 	for (MemoryContext cur = context; cur != NULL; cur = cur->parent)
 	{
-		MemoryStatsContextId *entry;
+		MemoryContextId *entry;
 		bool		found;
 
 		entry = hash_search(context_id_lookup, &cur, HASH_FIND, &found);
@@ -314,7 +312,7 @@ pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 	HTAB	   *context_id_lookup;
 
 	ctl.keysize = sizeof(MemoryContext);
-	ctl.entrysize = sizeof(MemoryStatsContextId);
+	ctl.entrysize = sizeof(MemoryContextId);
 	ctl.hcxt = CurrentMemoryContext;
 
 	context_id_lookup = hash_create("pg_get_backend_memory_contexts",
@@ -341,7 +339,7 @@ pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 
 	foreach_ptr(MemoryContextData, cur, contexts)
 	{
-		MemoryStatsContextId *entry;
+		MemoryContextId *entry;
 		bool		found;
 
 		/*
@@ -349,8 +347,8 @@ pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 		 * PutMemoryContextsStatsTupleStore needs this to populate the "path"
 		 * column with the parent context_ids.
 		 */
-		entry = (MemoryStatsContextId *) hash_search(context_id_lookup, &cur,
-													 HASH_ENTER, &found);
+		entry = (MemoryContextId *) hash_search(context_id_lookup, &cur,
+												HASH_ENTER, &found);
 		entry->context_id = context_id++;
 		Assert(!found);
 
@@ -437,10 +435,8 @@ pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
  *		wait for the results and display them.
  *
  * By default, only superusers or users with ROLE_PG_READ_ALL_STATS are allowed
- * to signal a process to return the memory contexts. This is because allowing
- * any users to issue this request at an unbounded rate would cause lots of
- * requests to be sent, which can lead to denial of service. Additional roles
- * can be permitted with GRANT.
+ * to signal a process to return the memory contexts.  Additional roles can be
+ * permitted with GRANT.
  *
  * On receipt of this signal, a backend or an auxiliary process sets the flag
  * in the signal handler, which causes the next CHECK_FOR_INTERRUPTS()
@@ -453,14 +449,14 @@ pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
  * at the end of the buffer.
  *
  * After sending the signal, wait on a condition variable. The publishing
- * backend, after copying the data to shared memory, sends signal on that
+ * backend, after copying the data to shared memory, sends a signal on that
  * condition variable. There is one condition variable per client process.
  * Once the condition variable is signalled, check if the latest memory context
  * information is available and display.
  *
  * If the publishing backend does not respond before the condition variable
- * times out, which is set to a predefined value MEMORY_STATS_MAX_TIMEOUT, give up
- * and return NULL.
+ * times out, which is set to a predefined value MEMORY_STATS_MAX_TIMEOUT, give
+ * up and return NULL.
  */
 Datum
 pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
@@ -495,12 +491,8 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 	 */
 	if (proc == NULL)
 	{
-		/*
-		 * This is a warning because we don't want to break loops.
-		 */
 		ereport(WARNING,
-				errmsg("PID %d is not a PostgreSQL server process",
-					   pid));
+				errmsg("PID %d is not a PostgreSQL server process", pid));
 		PG_RETURN_NULL();
 	}
 
@@ -508,7 +500,7 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 
 	/*
 	 * Create a DSA to allocate memory for copying memory contexts statistics.
-	 * Allocate the memory in the DSA and send dsa pointer to the server
+	 * Allocate the memory in the DSA and send DSA pointer to the server
 	 * process for storing the context statistics. If number of contexts
 	 * exceed a predefined limit (1MB), a cumulative total is stored for such
 	 * contexts.
@@ -521,8 +513,8 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 										 &found);
 
 	/*
-	 * The dsa pointers containing statistics for each client are stored in a
-	 * dshash table. In addition to dsa pointer, each entry in this table also
+	 * The DSA pointers containing statistics for each client are stored in a
+	 * dshash table. In addition to DSA pointer, each entry in this table also
 	 * contains information about the statistics, condition variable for
 	 * signalling between client and the server and miscellaneous data
 	 * specific to a request. There is one entry per client request in the
@@ -838,9 +830,9 @@ HandleGetMemoryContextInterrupt(void)
  * output.
  *
  * Statistics for all the processes are shared via the same dynamic shared
- * area. Individual statistics are tracked independently in
- * per-process DSA pointers. These pointers are stored in a dshash table with
- * key as requesting clients ProcNumber.
+ * area. Individual statistics are tracked independently in per-process DSA
+ * pointers. These pointers are stored in a dshash table with key as requesting
+ * clients ProcNumber.
  *
  * We calculate maximum number of context's statistics that can be displayed
  * using a pre-determined limit for memory available per process for this
@@ -848,11 +840,11 @@ HandleGetMemoryContextInterrupt(void)
  * context statistics if any are captured as a cumulative total at the end of
  * individual context's statistics.
  *
- * If summary is true, we capture the level 1 and level 2 contexts
- * statistics.  For that we traverse the memory context tree recursively in
- * depth first search manner to cover all the children of a parent context, to
- * be able to display a cumulative total of memory consumption by a parent at
- * level 2 and all its children.
+ * If summary is true, we capture the level 1 and level 2 contexts statistics.
+ * For that we traverse the memory context tree recursively in depth first
+ * search manner to cover all the children of a parent context, to be able to
+ * display a cumulative total of memory consumption by a parent at level 2 and
+ * all its children.
  */
 void
 ProcessGetMemoryContextInterrupt(void)
@@ -899,7 +891,7 @@ ProcessGetMemoryContextInterrupt(void)
 	 * similar to its local backend counterpart.
 	 */
 	ctl.keysize = sizeof(MemoryContext);
-	ctl.entrysize = sizeof(MemoryStatsContextId);
+	ctl.entrysize = sizeof(MemoryContextId);
 	ctl.hcxt = CurrentMemoryContext;
 
 	context_id_lookup = hash_create("pg_get_remote_backend_memory_contexts",
@@ -912,7 +904,7 @@ ProcessGetMemoryContextInterrupt(void)
 
 	/*
 	 * If DSA exists, created by another process requesting statistics, attach
-	 * to it. We expect the client process to create required DSA and Dshash
+	 * to it. We expect the client process to create required DSA and DSHash
 	 * table.
 	 */
 	if (MemoryStatsDsaArea == NULL)
@@ -923,7 +915,6 @@ ProcessGetMemoryContextInterrupt(void)
 		MemoryStatsDsHash = GetNamedDSHash("memory_context_statistics_dshash",
 										   &memctx_dsh_params, &found);
 
-
 	snprintf(key, CLIENT_KEY_SIZE, "%d", clientProcNumber);
 
 	/*
@@ -935,25 +926,24 @@ ProcessGetMemoryContextInterrupt(void)
 	entry = dshash_find_or_insert(MemoryStatsDsHash, key, &found);
 
 	/*
-	 * Entry has been deleted due to client process exit. Make sure that the
-	 * client always deletes the entry after taking required lock or this
-	 * function may end up writing to unallocated memory.
+	 * Check if the entry has been deleted due to calling process exiting, or
+	 * if the caller has timed out waiting for us and have issued a request to
+	 * another backend.
+	 *
+	 * XXX ?: Make sure that the client always deletes the entry after taking
+	 * required lock or this function may end up writing to unallocated
+	 * memory.
 	 */
-	if (!found)
+	if (!found || entry->target_server_id != MyProcPid)
 	{
 		entry->stats_written = false;
-		end_memorycontext_reporting(entry, oldcontext, context_id_lookup);
-		return;
-	}
 
-	/*
-	 * The client has timed out waiting for us to write statistics and is
-	 * requesting statistics from some other process
-	 */
-	if (entry->target_server_id != MyProcPid)
-	{
-		entry->stats_written = false;
-		end_memorycontext_reporting(entry, oldcontext, context_id_lookup);
+		dshash_release_lock(MemoryStatsDsHash, entry);
+
+		hash_destroy(context_id_lookup);
+		MemoryContextSwitchTo(oldcontext);
+		MemoryContextReset(memstats_ctx);
+
 		return;
 	}
 
@@ -966,7 +956,7 @@ ProcessGetMemoryContextInterrupt(void)
 	{
 		int			cxt_id = 0;
 		List	   *path = NIL;
-		MemoryStatsContextId *contextid_entry;
+		MemoryContextId *contextid_entry;
 
 		/* Copy TopMemoryContext statistics to DSA */
 		memset(&stat, 0, sizeof(stat));
@@ -976,9 +966,9 @@ ProcessGetMemoryContextInterrupt(void)
 		PublishMemoryContext(meminfo, cxt_id, TopMemoryContext, path, stat,
 							 1);
 
-		contextid_entry = (MemoryStatsContextId *) hash_search(context_id_lookup,
-															   &TopMemoryContext,
-															   HASH_ENTER, &found);
+		contextid_entry = (MemoryContextId *) hash_search(context_id_lookup,
+														  &TopMemoryContext,
+														  HASH_ENTER, &found);
 		Assert(!found);
 
 		/*
@@ -1001,8 +991,8 @@ ProcessGetMemoryContextInterrupt(void)
 			memset(&grand_totals, 0, sizeof(grand_totals));
 
 			cxt_id++;
-			contextid_entry = (MemoryStatsContextId *) hash_search(context_id_lookup,
-																   &c, HASH_ENTER, &found);
+			contextid_entry = (MemoryContextId *) hash_search(context_id_lookup,
+															  &c, HASH_ENTER, &found);
 			Assert(!found);
 			contextid_entry->context_id = cxt_id + 1;
 
@@ -1014,119 +1004,111 @@ ProcessGetMemoryContextInterrupt(void)
 								 grand_totals, num_contexts);
 		}
 		entry->total_stats = cxt_id + 1;
-
-		entry->stats_written = true;
-		end_memorycontext_reporting(entry, oldcontext, context_id_lookup);
-		/* Notify waiting client backend and return */
-		ConditionVariableSignal(&entry->memcxt_cv);
-		return;
 	}
-	foreach_ptr(MemoryContextData, cur, contexts)
+	else
 	{
-		List	   *path = NIL;
-		MemoryStatsContextId *contextid_entry;
+		foreach_ptr(MemoryContextData, cur, contexts)
+		{
+			List	   *path = NIL;
+			MemoryContextId *contextid_entry;
 
-		contextid_entry = (MemoryStatsContextId *) hash_search(context_id_lookup,
-															   &cur,
-															   HASH_ENTER, &found);
-		Assert(!found);
+			contextid_entry = (MemoryContextId *) hash_search(context_id_lookup,
+															  &cur,
+															  HASH_ENTER, &found);
+			Assert(!found);
 
-		/*
-		 * context id starts with 1
-		 */
-		contextid_entry->context_id = context_id + 1;
+			/*
+			 * context id starts with 1
+			 */
+			contextid_entry->context_id = context_id + 1;
+
+			/*
+			 * Figure out the transient context_id of this context and each of
+			 * its ancestors, to compute a path for this context.
+			 */
+			path = compute_context_path(cur, context_id_lookup);
+
+			/* Examine the context stats */
+			memset(&stat, 0, sizeof(stat));
+			(*cur->methods->stats) (cur, NULL, NULL, &stat, true);
+
+			/* Account for saving one statistics slot for cumulative reporting */
+			if (context_id < (MAX_MEMORY_CONTEXT_STATS_NUM - 1))
+			{
+				/* Copy statistics to DSA memory */
+				PublishMemoryContext(meminfo, context_id, cur, path, stat, 1);
+			}
+			else
+			{
+				meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].totalspace += stat.totalspace;
+				meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].nblocks += stat.nblocks;
+				meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].freespace += stat.freespace;
+				meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].freechunks += stat.freechunks;
+			}
+
+			/*
+			 * DSA max limit per process is reached, write aggregate of the
+			 * remaining statistics.
+			 *
+			 * We can store contexts from 0 to max_stats - 1. When context_id
+			 * is greater than max_stats, we stop reporting individual
+			 * statistics when context_id equals max_stats - 2. As we use
+			 * max_stats - 1 array slot for reporting cumulative statistics or
+			 * "Remaining Totals".
+			 */
+			if (context_id == (MAX_MEMORY_CONTEXT_STATS_NUM - 2))
+			{
+				int			namelen = strlen("Remaining Totals");
+
+				num_individual_stats = context_id + 1;
+				strlcpy(meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].name,
+						"Remaining Totals", namelen + 1);
+				meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].ident[0] = '\0';
+				meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].path[0] = 0;
+				meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].type = 0;
+			}
+			context_id++;
+
+			for (MemoryContext c = cur->firstchild; c != NULL; c = c->nextchild)
+				contexts = lappend(contexts, c);
+		}
 
 		/*
-		 * Figure out the transient context_id of this context and each of its
-		 * ancestors, to compute a path for this context.
+		 * Check if there are aggregated statistics or not in the result set.
+		 * Statistics are individually reported when context_id <= max_stats,
+		 * only if context_id > max_stats will there be aggregates.
 		 */
-		path = compute_context_path(cur, context_id_lookup);
-
-		/* Examine the context stats */
-		memset(&stat, 0, sizeof(stat));
-		(*cur->methods->stats) (cur, NULL, NULL, &stat, true);
-
-		/* Account for saving one statistics slot for cumulative reporting */
-		if (context_id < (MAX_MEMORY_CONTEXT_STATS_NUM - 1))
+		if (context_id <= MAX_MEMORY_CONTEXT_STATS_NUM)
 		{
-			/* Copy statistics to DSA memory */
-			PublishMemoryContext(meminfo, context_id, cur, path, stat, 1);
-		}
-		else
-		{
-			meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].totalspace += stat.totalspace;
-			meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].nblocks += stat.nblocks;
-			meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].freespace += stat.freespace;
-			meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].freechunks += stat.freechunks;
+			entry->total_stats = context_id;
+			meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].num_agg_stats = 1;
 		}
 
 		/*
-		 * DSA max limit per process is reached, write aggregate of the
-		 * remaining statistics.
-		 *
-		 * We can store contexts from 0 to max_stats - 1. When context_id is
-		 * greater than max_stats, we stop reporting individual statistics
-		 * when context_id equals max_stats - 2. As we use max_stats - 1 array
-		 * slot for reporting cumulative statistics or "Remaining Totals".
+		 * The number of contexts exceeded the space available, so report the
+		 * number of aggregated memory contexts
 		 */
-		if (context_id == (MAX_MEMORY_CONTEXT_STATS_NUM - 2))
+		else
 		{
-			int			namelen = strlen("Remaining Totals");
-
-			num_individual_stats = context_id + 1;
-			strlcpy(meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].name,
-					"Remaining Totals", namelen + 1);
-			meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].ident[0] = '\0';
-			meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].path[0] = 0;
-			meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].type = 0;
+			meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].num_agg_stats =
+				context_id - num_individual_stats;
+
+			/*
+			 * Total stats equals num_individual_stats + 1 record for
+			 * cumulative statistics.
+			 */
+			entry->total_stats = num_individual_stats + 1;
 		}
-		context_id++;
-
-		for (MemoryContext c = cur->firstchild; c != NULL; c = c->nextchild)
-			contexts = lappend(contexts, c);
 	}
 
-	/*
-	 * Statistics are not aggregated, i.e individual statistics reported when
-	 * context_id <= max_stats.
-	 */
-	if (context_id <= MAX_MEMORY_CONTEXT_STATS_NUM)
-	{
-		entry->total_stats = context_id;
-		meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].num_agg_stats = 1;
-	}
-	/* Report number of aggregated memory contexts */
-	else
-	{
-		meminfo[MAX_MEMORY_CONTEXT_STATS_NUM - 1].num_agg_stats = context_id
-			- num_individual_stats;
-
-		/*
-		 * Total stats equals num_individual_stats + 1 record for cumulative
-		 * statistics.
-		 */
-		entry->total_stats = num_individual_stats + 1;
-	}
 	entry->stats_written = true;
-	end_memorycontext_reporting(entry, oldcontext, context_id_lookup);
-	/* Notify waiting client backend and return */
-	ConditionVariableSignal(&entry->memcxt_cv);
-}
-
-/*
- * Clean up before exit from ProcessGetMemoryContextInterrupt
- */
-static void
-end_memorycontext_reporting(MemoryStatsDSHashEntry *entry,
-							MemoryContext oldcontext, HTAB *context_id_lookup)
-{
-	MemoryContext curr_ctx = CurrentMemoryContext;
-
 	dshash_release_lock(MemoryStatsDsHash, entry);
-
 	hash_destroy(context_id_lookup);
+
 	MemoryContextSwitchTo(oldcontext);
-	MemoryContextReset(curr_ctx);
+	MemoryContextReset(memstats_ctx);
+	/* Notify waiting client backend and return */
+	ConditionVariableSignal(&entry->memcxt_cv);
 }
 
 /*
@@ -1144,7 +1126,7 @@ compute_context_path(MemoryContext c, HTAB *context_id_lookup)
 
 	for (cur_context = c; cur_context != NULL; cur_context = cur_context->parent)
 	{
-		MemoryStatsContextId *cur_entry;
+		MemoryContextId *cur_entry;
 
 		cur_entry = hash_search(context_id_lookup, &cur_context, HASH_FIND, &found);
 
@@ -1167,8 +1149,8 @@ PublishMemoryContext(MemoryStatsEntry *memcxt_info, int curr_id,
 					 MemoryContext context, List *path,
 					 MemoryContextCounters stat, int num_contexts)
 {
-	const char *ident = context->ident;
-	const char *name = context->name;
+	char	   *ident = unconstify(char *, context->ident);
+	char	   *name = unconstify(char *, context->name);
 
 	/*
 	 * To be consistent with logging output, we label dynahash contexts with
@@ -1176,7 +1158,7 @@ PublishMemoryContext(MemoryStatsEntry *memcxt_info, int curr_id,
 	 */
 	if (context->ident && strncmp(context->name, "dynahash", 8) == 0)
 	{
-		name = context->ident;
+		name = unconstify(char *, context->ident);
 		ident = NULL;
 	}
 
@@ -1184,7 +1166,7 @@ PublishMemoryContext(MemoryStatsEntry *memcxt_info, int curr_id,
 	{
 		int			namelen = strlen(name);
 
-		if (strlen(name) >= MEMORY_CONTEXT_NAME_SHMEM_SIZE)
+		if (namelen >= MEMORY_CONTEXT_NAME_SHMEM_SIZE)
 			namelen = pg_mbcliplen(name, namelen,
 								   MEMORY_CONTEXT_NAME_SHMEM_SIZE - 1);
 
@@ -1212,7 +1194,7 @@ PublishMemoryContext(MemoryStatsEntry *memcxt_info, int curr_id,
 	else
 		memcxt_info[curr_id].ident[0] = '\0';
 
-	/* Allocate DSA memory for storing path information */
+	/* Store the path */
 	if (path == NIL)
 		memcxt_info[curr_id].path[0] = 0;
 	else
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index 4296667cbf0..617de0ebf91 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -18,7 +18,6 @@
 #define MEMUTILS_H
 
 #include "nodes/memnodes.h"
-#include "utils/dsa.h"
 
 /*
  * MaxAllocSize, MaxAllocHugeSize
@@ -48,6 +47,7 @@
 
 #define AllocHugeSizeIsValid(size)	((Size) (size) <= MaxAllocHugeSize)
 
+
 /*
  * Standard top-level memory contexts.
  *
diff --git a/src/test/modules/test_memcontext_reporting/t/001_memcontext_inj.pl b/src/test/modules/test_memcontext_reporting/t/001_memcontext_inj.pl
index 69d8489eb37..8fa12d1f693 100644
--- a/src/test/modules/test_memcontext_reporting/t/001_memcontext_inj.pl
+++ b/src/test/modules/test_memcontext_reporting/t/001_memcontext_inj.pl
@@ -11,7 +11,7 @@ use Test::More;
 
 if ($ENV{enable_injection_points} ne 'yes')
 {
-       plan skip_all => 'Injection points not supported by this build';
+	plan skip_all => 'Injection points not supported by this build';
 }
 my $psql_err;
 # Create and start a cluster with one node
@@ -21,8 +21,8 @@ $node->init(allows_streaming => 1);
 # and log_statement is dialled down since it otherwise will generate enormous
 # amounts of logging. Page verification failures are still logged.
 $node->append_conf(
-       'postgresql.conf',
-       qq[
+	'postgresql.conf',
+	qq[
 max_connections = 100
 log_statement = none
 ]);
@@ -30,29 +30,45 @@ $node->start;
 $node->safe_psql('postgres', 'CREATE EXTENSION test_memcontext_reporting;');
 $node->safe_psql('postgres', 'CREATE EXTENSION injection_points;');
 # Attaching to a client process injection point that throws an error
-$node->safe_psql('postgres', "select injection_points_attach('memcontext-client-crash', 'error');");
+$node->safe_psql('postgres',
+	"select injection_points_attach('memcontext-client-crash', 'error');");
 
-my $pid = $node->safe_psql('postgres', "SELECT pid from pg_stat_activity where backend_type='checkpointer'");
+my $pid = $node->safe_psql('postgres',
+	"SELECT pid from pg_stat_activity where backend_type='checkpointer'");
 print "PID";
 print $pid;
 
 #Client should have thrown error
-$node->psql('postgres', qq(select pg_get_process_memory_contexts($pid, true);), stderr => \$psql_err);
-like ( $psql_err, qr/error triggered for injection point memcontext-client-crash/);
+$node->psql(
+	'postgres',
+	qq(select pg_get_process_memory_contexts($pid, true);),
+	stderr => \$psql_err);
+like($psql_err,
+	qr/error triggered for injection point memcontext-client-crash/);
 
 #Query the same process after detaching the injection point, using some other client and it should succeed.
-$node->safe_psql('postgres', "select injection_points_detach('memcontext-client-crash');");
-my $topcontext_name = $node->safe_psql('postgres', "select name from pg_get_process_memory_contexts($pid, true) where path = '{1}';");
+$node->safe_psql('postgres',
+	"select injection_points_detach('memcontext-client-crash');");
+my $topcontext_name = $node->safe_psql('postgres',
+	"select name from pg_get_process_memory_contexts($pid, true) where path = '{1}';"
+);
 ok($topcontext_name = 'TopMemoryContext');
 
 # Attaching to a target process injection point that throws an error
-$node->safe_psql('postgres', "select injection_points_attach('memcontext-server-crash', 'error');");
+$node->safe_psql('postgres',
+	"select injection_points_attach('memcontext-server-crash', 'error');");
 
 #Server should have thrown error
-$node->psql('postgres', qq(select pg_get_process_memory_contexts($pid, true);), stderr => \$psql_err);
+$node->psql(
+	'postgres',
+	qq(select pg_get_process_memory_contexts($pid, true);),
+	stderr => \$psql_err);
 
 #Query the same process after detaching the injection point, using some other client and it should succeed.
-$node->safe_psql('postgres', "select injection_points_detach('memcontext-server-crash');");
-$topcontext_name = $node->safe_psql('postgres', "select name from pg_get_process_memory_contexts($pid, true) where path = '{1}';");
+$node->safe_psql('postgres',
+	"select injection_points_detach('memcontext-server-crash');");
+$topcontext_name = $node->safe_psql('postgres',
+	"select name from pg_get_process_memory_contexts($pid, true) where path = '{1}';"
+);
 ok($topcontext_name = 'TopMemoryContext');
 done_testing();
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index aa898f025c0..5e3122a468b 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -1686,9 +1686,9 @@ MemoryContextCallback
 MemoryContextCallbackFunction
 MemoryContextCounters
 MemoryContextData
+MemoryContextId
 MemoryContextMethodID
 MemoryContextMethods
-MemoryStatsContextId
 MemoryStatsEntry
 MemoryStatsDSHashEntry
 MemoryStatsPrintFunc
-- 
2.39.3 (Apple Git-146)

