From c4ea611583dd36c1a3facf7d3185c1e9e93b17b2 Mon Sep 17 00:00:00 2001
From: Rahila Syed <rahilasyed.90@gmail.com>
Date: Mon, 30 Jun 2025 12:11:00 +0530
Subject: [PATCH] Add pg_get_process_memory_context function

This adds a function for retrieving memory context statistics
and information from backends as well as auxiliary processes.
The intended usecase is cluster debugging when under memory
pressure or unanticipated memory usage characteristics.

When calling the function it sends a signal to the specified
process to submit statistics regarding its memory contexts
into dynamic shared memory.  Each memory context is returned
in detail, followed by a cumulative total in case the number
of contexts exceed the max allocated amount of shared memory.
Each process is limited to use at most 1Mb memory for this.

A summary can also be explicitly requested by the user, this
will return the TopMemoryContext and a cumulative total of
all lower contexts.

In order to not block on busy processes the caller specifies
the number of seconds during which to retry before timing out.
In the case where no statistics are published within the set
timeout, NULL is returned.
---
 doc/src/sgml/func.sgml                        | 164 ++++
 src/backend/catalog/system_views.sql          |   5 +
 src/backend/postmaster/autovacuum.c           |   4 +
 src/backend/postmaster/checkpointer.c         |   4 +
 src/backend/postmaster/interrupt.c            |   4 +
 src/backend/postmaster/pgarch.c               |   4 +
 src/backend/postmaster/startup.c              |   4 +
 src/backend/postmaster/walsummarizer.c        |   4 +
 src/backend/storage/ipc/ipci.c                |   3 +
 src/backend/storage/ipc/procsignal.c          |   3 +
 src/backend/storage/lmgr/lwlock.c             |   1 +
 src/backend/storage/lmgr/proc.c               |   1 +
 src/backend/tcop/postgres.c                   |   3 +
 .../utils/activity/wait_event_names.txt       |   1 +
 src/backend/utils/adt/mcxtfuncs.c             | 837 +++++++++++++++++-
 src/backend/utils/adt/pg_locale.c             |   1 -
 src/backend/utils/init/globals.c              |   1 +
 src/backend/utils/init/postinit.c             |   7 +
 src/backend/utils/mb/mbutils.c                |   1 -
 src/backend/utils/mmgr/mcxt.c                 |  71 +-
 src/include/catalog/pg_proc.dat               |  10 +
 src/include/miscadmin.h                       |   1 +
 src/include/storage/lwlock.h                  |   1 +
 src/include/storage/procsignal.h              |   1 +
 src/include/utils/memutils.h                  |  92 +-
 src/test/regress/expected/sysviews.out        |  19 +
 src/test/regress/sql/sysviews.sql             |  18 +
 src/tools/pgindent/typedefs.list              |   3 +
 28 files changed, 1227 insertions(+), 41 deletions(-)

diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
index c28aa71f570..ba082030ea2 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -28716,6 +28716,137 @@ acl      | {postgres=arwdDxtm/postgres,foo=r/postgres}
        </para></entry>
       </row>
 
+      <row>
+       <entry role="func_table_entry"><para role="func_signature">
+        <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>, <parameter>timeout</parameter> <type>float</type> )
+        <returnvalue>setof record</returnvalue>
+        ( <parameter>name</parameter> <type>text</type>,
+        <parameter>ident</parameter> <type>text</type>,
+        <parameter>type</parameter> <type>text</type>,
+        <parameter>path</parameter> <type>integer[]</type>,
+        <parameter>level</parameter> <type>integer</type>,
+        <parameter>total_bytes</parameter> <type>bigint</type>,
+        <parameter>total_nblocks</parameter> <type>bigint</type>,
+        <parameter>free_bytes</parameter> <type>bigint</type>,
+        <parameter>free_chunks</parameter> <type>bigint</type>,
+        <parameter>used_bytes</parameter> <type>bigint</type>,
+        <parameter>num_agg_contexts</parameter> <type>integer</type>,
+        <parameter>stats_timestamp</parameter> <type>timestamptz</type> )
+       </para>
+       <para>
+        This function handles requests to display the memory contexts of a
+        <productname>PostgreSQL</productname> process with the specified
+        process ID.  The function can be used to send requests to backends as
+        well as <glossterm linkend="glossary-auxiliary-proc">auxiliary processes</glossterm>.
+       </para>
+       <para>
+        The returned record contains extended statistics per each memory
+        context:
+        <itemizedlist spacing="compact">
+         <listitem>
+          <para>
+           <parameter>name</parameter> - The name of the memory context.
+          </para>
+         </listitem>
+         <listitem>
+          <para>
+           <parameter>ident</parameter> - Memory context ID (if any).
+          </para>
+         </listitem>
+         <listitem>
+          <para>
+           <parameter>type</parameter> - The type of memory context, possible
+           values are: AllocSet, Generation, Slab and Bump.
+          </para>
+         </listitem>
+         <listitem>
+          <para>
+           <parameter>path</parameter> - Memory contexts are organized in a
+           tree model with TopMemoryContext as the root, and all other memory
+           contexts as nodes in the tree. The <parameter>path</parameter>
+           displays the path from the root to the current memory context. The
+           path is limited to 100 children per node, which each node limited
+           to a max depth of 100, to preserve memory during reporting. The
+           printed path will also be limited to 100 nodes counting from the
+           TopMemoryContext.
+          </para>
+         </listitem>
+         <listitem>
+          <para>
+           <parameter>level</parameter> - The level in the tree of the current
+           memory context.
+          </para>
+         </listitem>
+         <listitem>
+          <para>
+           <parameter>total_bytes</parameter> - The total number of bytes
+           allocated to this memory context.
+          </para>
+         </listitem>
+         <listitem>
+          <para>
+           <parameter>total_nblocks</parameter> - The total number of blocks
+           used for the allocated memory.
+          </para>
+         </listitem>
+         <listitem>
+          <para>
+           <parameter>free_bytes</parameter> - The amount of free memory in
+           this memory context.
+          </para>
+         </listitem>
+         <listitem>
+          <para>
+           <parameter>free_chunks</parameter> - The number of chunks that
+           <parameter>free_bytes</parameter> corresponds to.
+          </para>
+         </listitem>
+         <listitem>
+          <para>
+           <parameter>used_bytes</parameter> - The total number of bytes
+           currently occupied.
+          </para>
+         </listitem>
+         <listitem>
+          <para>
+           <parameter>num_agg_contexts</parameter> - The number of memory
+           contexts aggregated in the displayed statistics.
+          </para>
+         </listitem>
+        </itemizedlist>
+       </para>
+       <para>
+        When <parameter>summary</parameter> is <literal>true</literal>, statistics
+        for memory contexts at levels 1 and 2 are displayed, with level 1
+        representing the root node (i.e., <literal>TopMemoryContext</literal>).
+        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>,
+        <literal>the num_agg_contexts</literal> value is <literal>1</literal>,
+        indicating that individual statistics are being displayed.
+       </para>
+       <para>
+        Busy processes can delay reporting memory context statistics,
+        <parameter>timeout</parameter> specifies the number of seconds
+        to wait for updated statistics. <parameter>timeout</parameter> can be
+        specified in fractions of a second.
+       </para>
+       <para>
+        After receiving memory context statistics from the target process, it
+        returns the results as one row per context.  If all the contexts don't
+        fit within the pre-determined size limit, the remaining context
+        statistics are aggregated and a cumulative total is displayed.  The
+        <literal>num_agg_contexts</literal> column indicates the number of
+        contexts aggregated in the displayed statistics.  When
+        <literal>num_agg_contexts</literal> is <literal>1</literal> it means
+        that the context statistics are displayed separately.
+       </para></entry>
+      </row>
+
       <row>
        <entry role="func_table_entry"><para role="func_signature">
         <indexterm>
@@ -28855,6 +28986,39 @@ LOG:  Grand total: 1651920 bytes in 201 blocks; 622360 free (88 chunks); 1029560
     because it may generate a large number of log messages.
    </para>
 
+   <para>
+    <function>pg_get_process_memory_contexts</function> can be used to request
+    memory contexts statistics of any <productname>PostgreSQL</productname>
+    process.  For example:
+<programlisting>
+postgres=# SELECT * FROM pg_get_process_memory_contexts(
+  (SELECT pid FROM pg_stat_activity
+    WHERE backend_type = 'checkpointer'),
+  false, 0.5) LIMIT 1;
+-[ RECORD 1 ]----+------------------------------
+name             | TopMemoryContext
+ident            |
+type             | AllocSet
+path             | {1}
+level            | 1
+total_bytes      | 90304
+total_nblocks    | 3
+free_bytes       | 2880
+free_chunks      | 1
+used_bytes       | 87424
+num_agg_contexts | 1
+</programlisting>
+    <note>
+     <para>
+      While <function>pg_get_process_memory_contexts</function> can be used to
+      query memory contexts of the local backend,
+      <structname>pg_backend_memory_contexts</structname>
+      (see <xref linkend="view-pg-backend-memory-contexts"/> for more details)
+      will be less resource intensive when only the local backend is of interest.
+     </para>
+    </note>
+   </para>
+
   </sect2>
 
   <sect2 id="functions-admin-backup">
diff --git a/src/backend/catalog/system_views.sql b/src/backend/catalog/system_views.sql
index b2d5332effc..33b5fcb9119 100644
--- a/src/backend/catalog/system_views.sql
+++ b/src/backend/catalog/system_views.sql
@@ -682,6 +682,11 @@ 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, float) FROM PUBLIC;
+GRANT EXECUTE ON FUNCTION
+	pg_get_process_memory_contexts(integer, boolean, float) TO pg_read_all_stats;
+
 -- Statistics views
 
 CREATE VIEW pg_stat_all_tables AS
diff --git a/src/backend/postmaster/autovacuum.c b/src/backend/postmaster/autovacuum.c
index 9474095f271..5e7e8081c05 100644
--- a/src/backend/postmaster/autovacuum.c
+++ b/src/backend/postmaster/autovacuum.c
@@ -781,6 +781,10 @@ ProcessAutoVacLauncherInterrupts(void)
 	if (LogMemoryContextPending)
 		ProcessLogMemoryContextInterrupt();
 
+	/* Publish memory contexts of this process */
+	if (PublishMemoryContextPending)
+		ProcessGetMemoryContextInterrupt();
+
 	/* Process sinval catchup interrupts that happened while sleeping */
 	ProcessCatchupInterrupt();
 }
diff --git a/src/backend/postmaster/checkpointer.c b/src/backend/postmaster/checkpointer.c
index fda91ffd1ce..d3cb3f1891c 100644
--- a/src/backend/postmaster/checkpointer.c
+++ b/src/backend/postmaster/checkpointer.c
@@ -663,6 +663,10 @@ ProcessCheckpointerInterrupts(void)
 	/* Perform logging of memory contexts of this process */
 	if (LogMemoryContextPending)
 		ProcessLogMemoryContextInterrupt();
+
+	/* Publish memory contexts of this process */
+	if (PublishMemoryContextPending)
+		ProcessGetMemoryContextInterrupt();
 }
 
 /*
diff --git a/src/backend/postmaster/interrupt.c b/src/backend/postmaster/interrupt.c
index 0ae9bf906ec..f24f574e748 100644
--- a/src/backend/postmaster/interrupt.c
+++ b/src/backend/postmaster/interrupt.c
@@ -48,6 +48,10 @@ ProcessMainLoopInterrupts(void)
 	/* Perform logging of memory contexts of this process */
 	if (LogMemoryContextPending)
 		ProcessLogMemoryContextInterrupt();
+
+	/* Publish memory contexts of this process */
+	if (PublishMemoryContextPending)
+		ProcessGetMemoryContextInterrupt();
 }
 
 /*
diff --git a/src/backend/postmaster/pgarch.c b/src/backend/postmaster/pgarch.c
index 78e39e5f866..ac97a39447c 100644
--- a/src/backend/postmaster/pgarch.c
+++ b/src/backend/postmaster/pgarch.c
@@ -867,6 +867,10 @@ ProcessPgArchInterrupts(void)
 	if (LogMemoryContextPending)
 		ProcessLogMemoryContextInterrupt();
 
+	/* Publish memory contexts of this process */
+	if (PublishMemoryContextPending)
+		ProcessGetMemoryContextInterrupt();
+
 	if (ConfigReloadPending)
 	{
 		char	   *archiveLib = pstrdup(XLogArchiveLibrary);
diff --git a/src/backend/postmaster/startup.c b/src/backend/postmaster/startup.c
index 27e86cf393f..7149a67fcbc 100644
--- a/src/backend/postmaster/startup.c
+++ b/src/backend/postmaster/startup.c
@@ -192,6 +192,10 @@ ProcessStartupProcInterrupts(void)
 	/* Perform logging of memory contexts of this process */
 	if (LogMemoryContextPending)
 		ProcessLogMemoryContextInterrupt();
+
+	/* Publish memory contexts of this process */
+	if (PublishMemoryContextPending)
+		ProcessGetMemoryContextInterrupt();
 }
 
 
diff --git a/src/backend/postmaster/walsummarizer.c b/src/backend/postmaster/walsummarizer.c
index 777c9a8d555..5d14684f6b2 100644
--- a/src/backend/postmaster/walsummarizer.c
+++ b/src/backend/postmaster/walsummarizer.c
@@ -879,6 +879,10 @@ ProcessWalSummarizerInterrupts(void)
 	/* Perform logging of memory contexts of this process */
 	if (LogMemoryContextPending)
 		ProcessLogMemoryContextInterrupt();
+
+	/* Publish memory contexts of this process */
+	if (PublishMemoryContextPending)
+		ProcessGetMemoryContextInterrupt();
 }
 
 /*
diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index 2fa045e6b0f..fe3d32e40b0 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -51,6 +51,7 @@
 #include "storage/sinvaladt.h"
 #include "utils/guc.h"
 #include "utils/injection_point.h"
+#include "utils/memutils.h"
 
 /* GUCs */
 int			shared_memory_type = DEFAULT_SHARED_MEMORY_TYPE;
@@ -150,6 +151,7 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, InjectionPointShmemSize());
 	size = add_size(size, SlotSyncShmemSize());
 	size = add_size(size, AioShmemSize());
+	size = add_size(size, MemoryContextKeysShmemSize() + sizeof(LWLockPadded));
 
 	/* include additional requested shmem from preload libraries */
 	size = add_size(size, total_addin_request);
@@ -343,6 +345,7 @@ CreateOrAttachShmemStructs(void)
 	WaitEventCustomShmemInit();
 	InjectionPointShmemInit();
 	AioShmemInit();
+	MemoryContextKeysShmemInit();
 }
 
 /*
diff --git a/src/backend/storage/ipc/procsignal.c b/src/backend/storage/ipc/procsignal.c
index a9bb540b55a..ce69e26d720 100644
--- a/src/backend/storage/ipc/procsignal.c
+++ b/src/backend/storage/ipc/procsignal.c
@@ -691,6 +691,9 @@ procsignal_sigusr1_handler(SIGNAL_ARGS)
 	if (CheckProcSignal(PROCSIG_LOG_MEMORY_CONTEXT))
 		HandleLogMemoryContextInterrupt();
 
+	if (CheckProcSignal(PROCSIG_GET_MEMORY_CONTEXT))
+		HandleGetMemoryContextInterrupt();
+
 	if (CheckProcSignal(PROCSIG_PARALLEL_APPLY_MESSAGE))
 		HandleParallelApplyMessageInterrupt();
 
diff --git a/src/backend/storage/lmgr/lwlock.c b/src/backend/storage/lmgr/lwlock.c
index 46f44bc4511..a7b5ede2b12 100644
--- a/src/backend/storage/lmgr/lwlock.c
+++ b/src/backend/storage/lmgr/lwlock.c
@@ -178,6 +178,7 @@ static const char *const BuiltinTrancheNames[] = {
 	[LWTRANCHE_XACT_SLRU] = "XactSLRU",
 	[LWTRANCHE_PARALLEL_VACUUM_DSA] = "ParallelVacuumDSA",
 	[LWTRANCHE_AIO_URING_COMPLETION] = "AioUringCompletion",
+	[LWTRANCHE_MEMORY_CONTEXT_KEYS] = "MemoryContextReportingKeys",
 };
 
 StaticAssertDecl(lengthof(BuiltinTrancheNames) ==
diff --git a/src/backend/storage/lmgr/proc.c b/src/backend/storage/lmgr/proc.c
index e9ef0fbfe32..f194e6b3dcc 100644
--- a/src/backend/storage/lmgr/proc.c
+++ b/src/backend/storage/lmgr/proc.c
@@ -50,6 +50,7 @@
 #include "storage/procsignal.h"
 #include "storage/spin.h"
 #include "storage/standby.h"
+#include "utils/memutils.h"
 #include "utils/timeout.h"
 #include "utils/timestamp.h"
 
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index 2f8c3d5f918..83db8a20efb 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -3533,6 +3533,9 @@ ProcessInterrupts(void)
 	if (LogMemoryContextPending)
 		ProcessLogMemoryContextInterrupt();
 
+	if (PublishMemoryContextPending)
+		ProcessGetMemoryContextInterrupt();
+
 	if (ParallelApplyMessagePending)
 		ProcessParallelApplyMessages();
 }
diff --git a/src/backend/utils/activity/wait_event_names.txt b/src/backend/utils/activity/wait_event_names.txt
index 4da68312b5f..78b1fa5ca43 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -161,6 +161,7 @@ WAL_RECEIVER_EXIT	"Waiting for the WAL receiver to exit."
 WAL_RECEIVER_WAIT_START	"Waiting for startup process to send initial data for streaming replication."
 WAL_SUMMARY_READY	"Waiting for a new WAL summary to be generated."
 XACT_GROUP_UPDATE	"Waiting for the group leader to update transaction status at transaction end."
+MEM_CXT_PUBLISH	"Waiting for a process to publish memory information."
 
 ABI_compatibility:
 
diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
index fe6dce9cba3..b44bcb18118 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -15,13 +15,38 @@
 
 #include "postgres.h"
 
+#include "access/twophase.h"
+#include "catalog/pg_authid_d.h"
 #include "funcapi.h"
 #include "mb/pg_wchar.h"
+#include "miscadmin.h"
+#include "storage/dsm_registry.h"
 #include "storage/proc.h"
 #include "storage/procarray.h"
+#include "utils/acl.h"
 #include "utils/array.h"
 #include "utils/builtins.h"
 #include "utils/hsearch.h"
+#include "utils/memutils.h"
+#include "utils/wait_event_types.h"
+
+#define CLIENT_KEY_SIZE 64
+
+static LWLock *client_keys_lock = NULL;
+static int *client_keys = NULL;
+static dshash_table *MemoryStatsDsHash = NULL;
+static dsa_area *MemoryStatsDsaArea = NULL;
+
+static void dsa_cleanup(MemoryStatsDSHashEntry *entry);
+static const char *ContextTypeToString(NodeTag type);
+static void PublishMemoryContext(MemoryStatsEntry *memcxt_info,
+								 int curr_id, MemoryContext context,
+								 List *path,
+								 MemoryContextCounters stat,
+								 int num_contexts, int max_levels);
+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.
@@ -89,7 +114,7 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	 */
 	for (MemoryContext cur = context; cur != NULL; cur = cur->parent)
 	{
-		MemoryContextId *entry;
+		MemoryStatsContextId *entry;
 		bool		found;
 
 		entry = hash_search(context_id_lookup, &cur, HASH_FIND, &found);
@@ -143,24 +168,7 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	else
 		nulls[1] = true;
 
-	switch (context->type)
-	{
-		case T_AllocSetContext:
-			type = "AllocSet";
-			break;
-		case T_GenerationContext:
-			type = "Generation";
-			break;
-		case T_SlabContext:
-			type = "Slab";
-			break;
-		case T_BumpContext:
-			type = "Bump";
-			break;
-		default:
-			type = "???";
-			break;
-	}
+	type = ContextTypeToString(context->type);
 
 	values[2] = CStringGetTextDatum(type);
 	values[3] = Int32GetDatum(list_length(path));	/* level */
@@ -175,6 +183,38 @@ PutMemoryContextsStatsTupleStore(Tuplestorestate *tupstore,
 	list_free(path);
 }
 
+/*
+ * ContextTypeToString
+ *		Returns a textual representation of a context type
+ *
+ * This should cover the same types as MemoryContextIsValid.
+ */
+const char *
+ContextTypeToString(NodeTag type)
+{
+	const char *context_type;
+
+	switch (type)
+	{
+		case T_AllocSetContext:
+			context_type = "AllocSet";
+			break;
+		case T_GenerationContext:
+			context_type = "Generation";
+			break;
+		case T_SlabContext:
+			context_type = "Slab";
+			break;
+		case T_BumpContext:
+			context_type = "Bump";
+			break;
+		default:
+			context_type = "???";
+			break;
+	}
+	return context_type;
+}
+
 /*
  * pg_get_backend_memory_contexts
  *		SQL SRF showing backend memory context.
@@ -189,7 +229,7 @@ pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 	HTAB	   *context_id_lookup;
 
 	ctl.keysize = sizeof(MemoryContext);
-	ctl.entrysize = sizeof(MemoryContextId);
+	ctl.entrysize = sizeof(MemoryStatsContextId);
 	ctl.hcxt = CurrentMemoryContext;
 
 	context_id_lookup = hash_create("pg_get_backend_memory_contexts",
@@ -216,7 +256,7 @@ pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 
 	foreach_ptr(MemoryContextData, cur, contexts)
 	{
-		MemoryContextId *entry;
+		MemoryStatsContextId *entry;
 		bool		found;
 
 		/*
@@ -224,8 +264,8 @@ pg_get_backend_memory_contexts(PG_FUNCTION_ARGS)
 		 * 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 = (MemoryStatsContextId *) hash_search(context_id_lookup, &cur,
+													 HASH_ENTER, &found);
 		entry->context_id = context_id++;
 		Assert(!found);
 
@@ -305,3 +345,754 @@ pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
 
 	PG_RETURN_BOOL(true);
 }
+
+/*
+ * pg_get_process_memory_contexts
+ *		Signal a backend or an auxiliary process to send its memory contexts,
+ *		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.
+ *
+ * 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()
+ * or process-specific interrupt handler to copy the memory context details
+ * to a dynamic shared memory space.
+ *
+ * We have defined a limit on DSA memory that could be allocated per process -
+ * if the process has more memory contexts than what can fit in the allocated
+ * size, the excess contexts are summarized and represented as cumulative total
+ * 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
+ * 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 MEMSTATS_WAIT_TIMEOUT, retry given that there is
+ * time left within the timeout specified by the user, before giving up and
+ * returning previously published statistics, if any. If no previous statistics
+ * exist, return NULL.
+ */
+#define MEMSTATS_WAIT_TIMEOUT 100
+Datum
+pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
+{
+	int			pid = PG_GETARG_INT32(0);
+	bool		summary = PG_GETARG_BOOL(1);
+	double		timeout = PG_GETARG_FLOAT8(2);
+	PGPROC	   *proc;
+	ProcNumber	procNumber = INVALID_PROC_NUMBER;
+	bool		proc_is_aux = false;
+	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	MemoryStatsEntry *memcxt_info;
+	MemoryStatsDSHashEntry *entry;
+	bool		found;
+	char		key[CLIENT_KEY_SIZE];
+
+	/*
+	 * See if the process with given pid is a backend or an auxiliary process
+	 * and remember the type for when we requery the process later.
+	 */
+	proc = BackendPidGetProc(pid);
+	if (proc == NULL)
+	{
+		proc = AuxiliaryPidGetProc(pid);
+		proc_is_aux = true;
+	}
+
+	/*
+	 * BackendPidGetProc() and AuxiliaryPidGetProc() return NULL if the pid
+	 * isn't valid; this is however not a problem and leave with a WARNING.
+	 * See comment in pg_log_backend_memory_contexts for a discussion on this.
+	 */
+	if (proc == NULL)
+	{
+		/*
+		 * This is just a warning so a loop-through-resultset will not abort
+		 * if one backend terminated on its own during the run.
+		 */
+		ereport(WARNING,
+				errmsg("PID %d is not a PostgreSQL server process", pid));
+		PG_RETURN_NULL();
+	}
+
+	InitMaterializedSRF(fcinfo, 0);
+
+	procNumber = GetNumberFromPGProc(proc);
+
+	/*
+	 * Create a DSA to allocate memory for copying memory contexts statistics.
+	 * 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.
+	 *
+	 * The DSA is created once for the lifetime of the server, and only
+	 * attached in subsequent calls.
+	 */
+	if (MemoryStatsDsaArea == NULL)
+		MemoryStatsDsaArea = GetNamedDSA("memory_context_statistics_dsa", &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
+	 * 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
+	 * hash table.
+	 */
+	if (MemoryStatsDsHash == NULL)
+		MemoryStatsDsHash = GetNamedDSHash("memory_context_statistics_dshash", &memctx_dsh_params, &found);
+
+	snprintf(key, sizeof(key), "%d", MyProcNumber);
+
+	/*
+	 * Check if the publishing process slot is empty and store this clients
+	 * key i.e its procNumber. This informs the publishing process that it is
+	 * supposed to write statistics in the hash entry corresponding to this
+	 * client.
+	 */
+	LWLockAcquire(client_keys_lock, LW_EXCLUSIVE);
+	if (client_keys[procNumber] == -1)
+		client_keys[procNumber] = MyProcNumber;
+	else
+	{
+		ereport(WARNING,
+				errmsg("server process is processing previous request %d: %m", pid));
+		LWLockRelease(client_keys_lock);
+		PG_RETURN_NULL();
+	}
+	LWLockRelease(client_keys_lock);
+
+	/*
+	 * Insert an entry for this client in DSHASH table the first time this
+	 * function is called. This entry is deleted when the process exits in
+	 * before_shmem_exit call.
+	 *
+	 * dshash_find_or_insert locks the entry to prevent the publisher from
+	 * reading before client has updated the entry.
+	 */
+	entry = dshash_find_or_insert(MemoryStatsDsHash, key, &found);
+	if (!found)
+		ConditionVariableInit(&entry->memcxt_cv);
+
+	/*
+	 * Allocate 1MB of memory for the backend to publish its statistics on
+	 * every call to this function. The meomry is freed at the end of the
+	 * function.
+	 */
+	Assert(!DsaPointerIsValid(entry->memstats_dsa_pointer));
+	entry->memstats_dsa_pointer =
+		dsa_allocate0(MemoryStatsDsaArea, MEMORY_CONTEXT_REPORT_MAX_PER_BACKEND);
+	entry->summary = summary;
+	dshash_release_lock(MemoryStatsDsHash, entry);
+
+	/*
+	 * Send a signal to a PostgreSQL process, informing it we want it to
+	 * produce information about its memory contexts.
+	 */
+	if (SendProcSignal(pid, PROCSIG_GET_MEMORY_CONTEXT, procNumber) < 0)
+	{
+		ereport(WARNING,
+				errmsg("could not send signal to process %d: %m", pid));
+		PG_RETURN_NULL();
+	}
+
+	while (1)
+	{
+
+		entry = dshash_find_or_insert(MemoryStatsDsHash, key, &found);
+		Assert(found);
+
+		/*
+		 * We expect to come out of sleep when the requested process has
+		 * finished publishing the statistics, verified using the correct
+		 * entry in the proc_id field.
+		 *
+		 * Make sure that the information belongs to pid we requested
+		 * information for, Otherwise loop back and wait for the server
+		 * process to finish publishing statistics.
+		 *
+		 * Note in procnumber.h file says that a procNumber can be re-used for
+		 * a different backend immediately after a backend exits. In case an
+		 * old process' data was there and not updated by the current process
+		 * in the slot identified by the procNumber, the pid of the requested
+		 * process and the proc_id might not match.
+		 *
+		 */
+		if (entry->proc_id == pid)
+			break;
+
+		dshash_release_lock(MemoryStatsDsHash, entry);
+
+		/*
+		 * Recheck the state of the backend before sleeping on the condition
+		 * variable to ensure the process is still alive.  Only check the
+		 * relevant process type based on the earlier PID check.
+		 */
+		if (proc_is_aux)
+			proc = AuxiliaryPidGetProc(pid);
+		else
+			proc = BackendPidGetProc(pid);
+
+		/*
+		 * The process ending during memory context processing is not an
+		 * error.
+		 */
+		if (proc == NULL)
+		{
+			ereport(WARNING,
+					errmsg("PID %d is no longer a PostgreSQL server process",
+						   pid));
+			dsa_cleanup(entry);
+			PG_RETURN_NULL();
+		}
+
+
+		/*
+		 * Wait for the timeout as defined by the user. If no statistics are
+		 * available within the allowed time then return NULL. The timer is
+		 * defined in milliseconds since that's what the condition variable
+		 * sleep uses.
+		 */
+		if (ConditionVariableTimedSleep(&entry->memcxt_cv,
+										(timeout * 1000), WAIT_EVENT_MEM_CXT_PUBLISH))
+		{
+			/* Timeout has expired, return NULL */
+			dsa_cleanup(entry);
+			PG_RETURN_NULL();
+		}
+	}
+
+	/*
+	 * Backend has finished publishing the stats, project them.
+	 */
+	memcxt_info = (MemoryStatsEntry *)
+		dsa_get_address(MemoryStatsDsaArea, entry->memstats_dsa_pointer);
+
+#define PG_GET_PROCESS_MEMORY_CONTEXTS_COLS	11
+	for (int i = 0; i < entry->total_stats; i++)
+	{
+		ArrayType  *path_array;
+		int			path_length;
+		Datum		values[PG_GET_PROCESS_MEMORY_CONTEXTS_COLS];
+		bool		nulls[PG_GET_PROCESS_MEMORY_CONTEXTS_COLS];
+		Datum	   *path_datum = NULL;
+
+		memset(values, 0, sizeof(values));
+		memset(nulls, 0, sizeof(nulls));
+
+		if (memcxt_info[i].name[0] != '\0')
+		{
+			values[0] = CStringGetTextDatum(memcxt_info[i].name);
+		}
+		else
+			nulls[0] = true;
+
+		if (memcxt_info[i].ident[0] != '\0')
+		{
+			values[1] = CStringGetTextDatum(memcxt_info[i].ident);
+		}
+		else
+			nulls[1] = true;
+
+		values[2] = CStringGetTextDatum(ContextTypeToString(memcxt_info[i].type));
+
+		path_length = memcxt_info[i].path_length;
+		path_datum = (Datum *) palloc(path_length * sizeof(Datum));
+		if (memcxt_info[i].path[0] != 0)
+		{
+			for (int j = 0; j < path_length; j++)
+				path_datum[j] = Int32GetDatum(memcxt_info[i].path[j]);
+			path_array = construct_array_builtin(path_datum, path_length, INT4OID);
+			values[3] = PointerGetDatum(path_array);
+		}
+		else
+			nulls[3] = true;
+
+		values[4] = Int32GetDatum(memcxt_info[i].levels);
+		values[5] = Int64GetDatum(memcxt_info[i].totalspace);
+		values[6] = Int64GetDatum(memcxt_info[i].nblocks);
+		values[7] = Int64GetDatum(memcxt_info[i].freespace);
+		values[8] = Int64GetDatum(memcxt_info[i].freechunks);
+		values[9] = Int64GetDatum(memcxt_info[i].totalspace -
+								  memcxt_info[i].freespace);
+		values[10] = Int32GetDatum(memcxt_info[i].num_agg_stats);
+
+		tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc,
+							 values, nulls);
+	}
+	dsa_cleanup(entry);
+	dshash_release_lock(MemoryStatsDsHash, entry);
+
+	ConditionVariableCancelSleep();
+
+	PG_RETURN_NULL();
+}
+
+static void
+dsa_cleanup(MemoryStatsDSHashEntry *entry)
+{
+	Assert(MemoryStatsDsaArea != NULL);
+	dsa_free(MemoryStatsDsaArea, entry->memstats_dsa_pointer);
+	entry->memstats_dsa_pointer = InvalidDsaPointer;
+	entry->proc_id = 0;
+}
+void
+MemoryContextKeysShmemInit(void)
+{
+	bool		found;
+
+	client_keys = (int *)
+		ShmemInitStruct("MemoryContextKeys",
+						MemoryContextKeysShmemSize() + sizeof(LWLockPadded), &found);
+	client_keys_lock = (LWLock *) ((char *) client_keys + MemoryContextKeysShmemSize());
+
+	if (!found)
+	{
+		MemSet(client_keys, -1, MemoryContextKeysShmemSize());
+		LWLockInitialize(client_keys_lock, LWTRANCHE_MEMORY_CONTEXT_KEYS);
+	}
+}
+
+Size
+MemoryContextKeysShmemSize(void)
+{
+	Size		sz = 0;
+	Size		TotalProcs = 0;
+
+	TotalProcs = add_size(TotalProcs, NUM_AUXILIARY_PROCS);
+	TotalProcs = add_size(TotalProcs, MaxBackends);
+	sz = add_size(sz, mul_size(TotalProcs, sizeof(int)));
+
+	return sz;
+}
+
+/*
+ * HandleGetMemoryContextInterrupt
+ *		Handle receipt of an interrupt indicating a request to publish memory
+ *		contexts statistics.
+ *
+ * All the actual work is deferred to ProcessGetMemoryContextInterrupt() as
+ * this cannot be performed in a signal handler.
+ */
+void
+HandleGetMemoryContextInterrupt(void)
+{
+	InterruptPending = true;
+	PublishMemoryContextPending = true;
+	/* latch will be set by procsignal_sigusr1_handler */
+}
+
+/*
+ * ProcessGetMemoryContextInterrupt
+ *		Generate information about memory contexts used by the process.
+ *
+ * Performs a breadth first search on the memory context tree, thus parents
+ * statistics are reported before their children in the monitoring function
+ * output.
+ *
+ * Statistics for all the processes are shared via the same dynamic shared
+ * area.  Statistics written by each process are tracked independently in
+ * per-process DSA pointers. These pointers are stored in static shared memory.
+ *
+ * We calculate maximum number of context's statistics that can be displayed
+ * using a pre-determined limit for memory available per process for this
+ * utility maximum size of statistics for each context.  The remaining 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.
+ */
+void
+ProcessGetMemoryContextInterrupt(void)
+{
+	List	   *contexts;
+	HASHCTL		ctl;
+	HTAB	   *context_id_lookup;
+	int			context_id = 0;
+	MemoryStatsEntry *meminfo;
+	bool		summary = false;
+	MemoryContextCounters stat;
+	int			num_individual_stats = 0;
+	bool		found;
+	MemoryStatsDSHashEntry *entry;
+	char		key[CLIENT_KEY_SIZE];
+	int			clientProcNumber;
+	MemoryContext memstats_ctx = NULL;
+	MemoryContext oldcontext = NULL;
+
+	PublishMemoryContextPending = false;
+
+	/*
+	 * Create a new memory context which is not a part of TopMemoryContext
+	 * tree. This context is used to allocate all memory in this function.
+	 * This helps in keeping the memory allocation in this function to report
+	 * memory consumption statistics separate. So that it does not affect the
+	 * output of this function.
+	 */
+	memstats_ctx = AllocSetContextCreate((MemoryContext) NULL, "publish_memory_context_statistics",
+										 ALLOCSET_SMALL_SIZES);
+	oldcontext = MemoryContextSwitchTo(memstats_ctx);
+
+	/*
+	 * The hash table is used for constructing "path" column of the view,
+	 * similar to its local backend counterpart.
+	 */
+	ctl.keysize = sizeof(MemoryContext);
+	ctl.entrysize = sizeof(MemoryStatsContextId);
+	ctl.hcxt = CurrentMemoryContext;
+
+	context_id_lookup = hash_create("pg_get_remote_backend_memory_contexts",
+									256,
+									&ctl,
+									HASH_ELEM | HASH_BLOBS | HASH_CONTEXT);
+
+	/* List of contexts to process in the next round - start at the top. */
+	contexts = list_make1(TopMemoryContext);
+
+	/*
+	 * If DSA exists, created by another process requesting statistics, attach
+	 * to it. We expect the client process to create required DSA and Dshash
+	 * table.
+	 */
+	if (MemoryStatsDsaArea == NULL)
+		MemoryStatsDsaArea = GetNamedDSA("memory_context_statistics_dsa", &found);
+
+	if (MemoryStatsDsHash == NULL)
+		MemoryStatsDsHash = GetNamedDSHash("memory_context_statistics_dshash", &memctx_dsh_params, &found);
+
+	/* Retreive the client key fo publishing statistics */
+	LWLockAcquire(client_keys_lock, LW_EXCLUSIVE);
+	Assert(client_keys[MyProcNumber] != -1);
+	clientProcNumber = client_keys[MyProcNumber];
+	LWLockRelease(client_keys_lock);
+
+	snprintf(key, CLIENT_KEY_SIZE, "%d", clientProcNumber);
+
+	entry = dshash_find_or_insert(MemoryStatsDsHash, key, &found);
+	summary = entry->summary;
+
+	/*
+	 * The entry lock is held by dshash_find_or_insert to protect writes to
+	 * process specific memory. Two different processes publishing statistics
+	 * do not block each other.
+	 */
+	entry->proc_id = MyProcPid;
+
+	/* Should be allocated by a client backend that is requesting statistics */
+	Assert(entry->memstats_dsa_pointer != InvalidDsaPointer);
+	meminfo = (MemoryStatsEntry *)
+		dsa_get_address(MemoryStatsDsaArea, entry->memstats_dsa_pointer);
+
+	if (summary)
+	{
+		int			cxt_id = 0;
+		List	   *path = NIL;
+		MemoryStatsContextId *contextid_entry;
+
+		/* Copy TopMemoryContext statistics to DSA */
+		memset(&stat, 0, sizeof(stat));
+		(*TopMemoryContext->methods->stats) (TopMemoryContext, NULL, NULL,
+											 &stat, true);
+		path = lcons_int(1, path);
+		PublishMemoryContext(meminfo, cxt_id, TopMemoryContext, path, stat,
+							 1, 100);
+
+		contextid_entry = (MemoryStatsContextId *) hash_search(context_id_lookup, &TopMemoryContext,
+															   HASH_ENTER, &found);
+		Assert(!found);
+
+		/*
+		 * context id starts with 1
+		 */
+		contextid_entry->context_id = cxt_id + 1;
+
+		/*
+		 * Copy statistics for each of TopMemoryContexts children.  This
+		 * includes statistics of at most 100 children per node, with each
+		 * child node limited to a depth of 100 in its subtree.
+		 */
+		for (MemoryContext c = TopMemoryContext->firstchild; c != NULL;
+			 c = c->nextchild)
+		{
+			MemoryContextCounters grand_totals;
+			int			num_contexts = 0;
+
+			path = NIL;
+			memset(&grand_totals, 0, sizeof(grand_totals));
+
+			cxt_id++;
+			contextid_entry = (MemoryStatsContextId *) hash_search(context_id_lookup, &c,
+																   HASH_ENTER, &found);
+			Assert(!found);
+			contextid_entry->context_id = cxt_id + 1;
+
+			MemoryContextStatsInternal(c, 1, 100, 100, &grand_totals,
+									   PRINT_STATS_NONE, &num_contexts);
+
+			path = compute_context_path(c, context_id_lookup);
+
+			PublishMemoryContext(meminfo, cxt_id, c, path,
+								 grand_totals, num_contexts, 100);
+		}
+		entry->total_stats = cxt_id + 1;
+
+		/* Notify waiting backends and return */
+		end_memorycontext_reporting(entry, oldcontext, context_id_lookup);
+		return;
+	}
+	foreach_ptr(MemoryContextData, cur, contexts)
+	{
+		List	   *path = NIL;
+		MemoryStatsContextId *contextid_entry;
+
+		contextid_entry = (MemoryStatsContextId *) hash_search(context_id_lookup, &cur,
+															   HASH_ENTER, &found);
+		Assert(!found);
+
+		/*
+		 * 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, 100);
+		}
+		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);
+	}
+
+	/*
+	 * 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;
+	}
+
+	/* Notify waiting backends and return */
+	end_memorycontext_reporting(entry, oldcontext, context_id_lookup);
+}
+
+/*
+ * Update timestamp and signal all the waiting client backends after copying
+ * all the statistics.
+ */
+static void
+end_memorycontext_reporting(MemoryStatsDSHashEntry *entry, MemoryContext oldcontext, HTAB *context_id_lookup)
+{
+	MemoryContext curr_ctx = CurrentMemoryContext;
+
+	dshash_release_lock(MemoryStatsDsHash, entry);
+	ConditionVariableBroadcast(&entry->memcxt_cv);
+
+	/*
+	 * Empty this processes slot, so other clients can request memory
+	 * statistics
+	 */
+	LWLockAcquire(client_keys_lock, LW_EXCLUSIVE);
+	client_keys[MyProcNumber] = -1;
+	LWLockRelease(client_keys_lock);
+
+	hash_destroy(context_id_lookup);
+	MemoryContextSwitchTo(oldcontext);
+	MemoryContextReset(curr_ctx);
+}
+
+/*
+ * compute_context_path
+ *
+ * Append the transient context_id of this context and each of its ancestors
+ * to a list, in order to compute a path.
+ */
+static List *
+compute_context_path(MemoryContext c, HTAB *context_id_lookup)
+{
+	bool		found;
+	List	   *path = NIL;
+	MemoryContext cur_context;
+
+	for (cur_context = c; cur_context != NULL; cur_context = cur_context->parent)
+	{
+		MemoryStatsContextId *cur_entry;
+
+		cur_entry = hash_search(context_id_lookup, &cur_context, HASH_FIND, &found);
+
+		if (!found)
+			elog(ERROR, "hash table corrupted, can't construct path value");
+
+		path = lcons_int(cur_entry->context_id, path);
+	}
+
+	return path;
+}
+
+/*
+ * PublishMemoryContext
+ *
+ * Copy the memory context statistics of a single context to a DSA memory
+ */
+static void
+PublishMemoryContext(MemoryStatsEntry *memcxt_info, int curr_id,
+					 MemoryContext context, List *path,
+					 MemoryContextCounters stat, int num_contexts, int max_levels)
+{
+	const char *ident = context->ident;
+	const char *name = context->name;
+
+	/*
+	 * To be consistent with logging output, we label dynahash contexts with
+	 * just the hash table name as with MemoryContextStatsPrint().
+	 */
+	if (context->ident && strncmp(context->name, "dynahash", 8) == 0)
+	{
+		name = context->ident;
+		ident = NULL;
+	}
+
+	if (name != NULL)
+	{
+		int			namelen = strlen(name);
+
+		if (strlen(name) >= MEMORY_CONTEXT_IDENT_SHMEM_SIZE)
+			namelen = pg_mbcliplen(name, namelen,
+								   MEMORY_CONTEXT_IDENT_SHMEM_SIZE - 1);
+
+		strlcpy(memcxt_info[curr_id].name, name, namelen + 1);
+	}
+	else
+		/* Clearing the array */
+		memcxt_info[curr_id].name[0] = '\0';
+
+	/* Trim and copy the identifier if it is not set to NULL */
+	if (ident != NULL)
+	{
+		int			idlen = strlen(context->ident);
+
+		/*
+		 * Some identifiers such as SQL query string can be very long,
+		 * truncate oversize identifiers.
+		 */
+		if (idlen >= MEMORY_CONTEXT_IDENT_SHMEM_SIZE)
+			idlen = pg_mbcliplen(ident, idlen,
+								 MEMORY_CONTEXT_IDENT_SHMEM_SIZE - 1);
+
+		strlcpy(memcxt_info[curr_id].ident, ident, idlen + 1);
+	}
+	else
+		memcxt_info[curr_id].ident[0] = '\0';
+
+	/* Allocate DSA memory for storing path information */
+	if (path == NIL)
+		memcxt_info[curr_id].path[0] = 0;
+	else
+	{
+		int			levels = Min(list_length(path), max_levels);
+
+		memcxt_info[curr_id].path_length = levels;
+		memcxt_info[curr_id].levels = list_length(path);
+
+		foreach_int(i, path)
+		{
+			memcxt_info[curr_id].path[foreach_current_index(i)] = i;
+			if (--levels == 0)
+				break;
+		}
+	}
+	memcxt_info[curr_id].type = context->type;
+	memcxt_info[curr_id].totalspace = stat.totalspace;
+	memcxt_info[curr_id].nblocks = stat.nblocks;
+	memcxt_info[curr_id].freespace = stat.freespace;
+	memcxt_info[curr_id].freechunks = stat.freechunks;
+	memcxt_info[curr_id].num_agg_stats = num_contexts;
+}
+
+void
+AtProcExit_memstats_cleanup(int code, Datum arg)
+{
+	int			idx = MyProcNumber;
+	MemoryStatsDSHashEntry *entry;
+	char		key[CLIENT_KEY_SIZE];
+	bool		found;
+
+	if (MemoryStatsDsHash != NULL)
+	{
+		snprintf(key, CLIENT_KEY_SIZE, "%d", idx);
+		entry = dshash_find_or_insert(MemoryStatsDsHash, key, &found);
+
+		if (found)
+		{
+			if (MemoryStatsDsaArea != NULL &&
+				DsaPointerIsValid(entry->memstats_dsa_pointer))
+				dsa_free(MemoryStatsDsaArea, entry->memstats_dsa_pointer);
+		}
+		dshash_delete_entry(MemoryStatsDsHash, entry);
+	}
+}
diff --git a/src/backend/utils/adt/pg_locale.c b/src/backend/utils/adt/pg_locale.c
index 97c2ac1faf9..ab768a7a91f 100644
--- a/src/backend/utils/adt/pg_locale.c
+++ b/src/backend/utils/adt/pg_locale.c
@@ -45,7 +45,6 @@
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
 #include "utils/pg_locale.h"
-#include "utils/relcache.h"
 #include "utils/syscache.h"
 
 #ifdef WIN32
diff --git a/src/backend/utils/init/globals.c b/src/backend/utils/init/globals.c
index d31cb45a058..92b0446b80c 100644
--- a/src/backend/utils/init/globals.c
+++ b/src/backend/utils/init/globals.c
@@ -39,6 +39,7 @@ volatile sig_atomic_t TransactionTimeoutPending = false;
 volatile sig_atomic_t IdleSessionTimeoutPending = false;
 volatile sig_atomic_t ProcSignalBarrierPending = false;
 volatile sig_atomic_t LogMemoryContextPending = false;
+volatile sig_atomic_t PublishMemoryContextPending = false;
 volatile sig_atomic_t IdleStatsUpdateTimeoutPending = false;
 volatile uint32 InterruptHoldoffCount = 0;
 volatile uint32 QueryCancelHoldoffCount = 0;
diff --git a/src/backend/utils/init/postinit.c b/src/backend/utils/init/postinit.c
index c86ceefda94..89d72cdd5ff 100644
--- a/src/backend/utils/init/postinit.c
+++ b/src/backend/utils/init/postinit.c
@@ -663,6 +663,13 @@ BaseInit(void)
 	 * drop ephemeral slots, which in turn triggers stats reporting.
 	 */
 	ReplicationSlotInitialize();
+
+	/*
+	 * The before shmem exit callback frees the DSA memory occupied by the
+	 * latest memory context statistics that could be published by this proc
+	 * if requested.
+	 */
+	before_shmem_exit(AtProcExit_memstats_cleanup, 0);
 }
 
 
diff --git a/src/backend/utils/mb/mbutils.c b/src/backend/utils/mb/mbutils.c
index 886ecbad871..308016d7763 100644
--- a/src/backend/utils/mb/mbutils.c
+++ b/src/backend/utils/mb/mbutils.c
@@ -39,7 +39,6 @@
 #include "mb/pg_wchar.h"
 #include "utils/fmgrprotos.h"
 #include "utils/memutils.h"
-#include "utils/relcache.h"
 #include "varatt.h"
 
 /*
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 15fa4d0a55e..d6b081d1bbf 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -23,6 +23,7 @@
 
 #include "mb/pg_wchar.h"
 #include "miscadmin.h"
+#include "utils/hsearch.h"
 #include "utils/memdebug.h"
 #include "utils/memutils.h"
 #include "utils/memutils_internal.h"
@@ -159,10 +160,6 @@ MemoryContext PortalContext = NULL;
 
 static void MemoryContextDeleteOnly(MemoryContext context);
 static void MemoryContextCallResetCallbacks(MemoryContext context);
-static void MemoryContextStatsInternal(MemoryContext context, int level,
-									   int max_level, int max_children,
-									   MemoryContextCounters *totals,
-									   bool print_to_stderr);
 static void MemoryContextStatsPrint(MemoryContext context, void *passthru,
 									const char *stats_string,
 									bool print_to_stderr);
@@ -831,11 +828,19 @@ MemoryContextStatsDetail(MemoryContext context,
 						 bool print_to_stderr)
 {
 	MemoryContextCounters grand_totals;
+	int			num_contexts;
+	PrintDestination print_location;
 
 	memset(&grand_totals, 0, sizeof(grand_totals));
 
+	if (print_to_stderr)
+		print_location = PRINT_STATS_TO_STDERR;
+	else
+		print_location = PRINT_STATS_TO_LOGS;
+
+	/* num_contexts report number of contexts aggregated in the output */
 	MemoryContextStatsInternal(context, 1, max_level, max_children,
-							   &grand_totals, print_to_stderr);
+							   &grand_totals, print_location, &num_contexts);
 
 	if (print_to_stderr)
 		fprintf(stderr,
@@ -870,13 +875,14 @@ MemoryContextStatsDetail(MemoryContext context,
  *		One recursion level for MemoryContextStats
  *
  * Print stats for this context if possible, but in any case accumulate counts
- * into *totals (if not NULL).
+ * into *totals (if not NULL). The callers should make sure that print_location
+ * is set to PRINT_STATS_TO_STDERR or PRINT_STATS_TO_LOGS or PRINT_STATS_NONE.
  */
-static void
+void
 MemoryContextStatsInternal(MemoryContext context, int level,
 						   int max_level, int max_children,
 						   MemoryContextCounters *totals,
-						   bool print_to_stderr)
+						   PrintDestination print_location, int *num_contexts)
 {
 	MemoryContext child;
 	int			ichild;
@@ -884,10 +890,39 @@ MemoryContextStatsInternal(MemoryContext context, int level,
 	Assert(MemoryContextIsValid(context));
 
 	/* Examine the context itself */
-	context->methods->stats(context,
-							MemoryContextStatsPrint,
-							&level,
-							totals, print_to_stderr);
+	switch (print_location)
+	{
+		case PRINT_STATS_TO_STDERR:
+			context->methods->stats(context,
+									MemoryContextStatsPrint,
+									&level,
+									totals, true);
+			break;
+
+		case PRINT_STATS_TO_LOGS:
+			context->methods->stats(context,
+									MemoryContextStatsPrint,
+									&level,
+									totals, false);
+			break;
+
+		case PRINT_STATS_NONE:
+
+			/*
+			 * Do not print the statistics if print_location is
+			 * PRINT_STATS_NONE, only compute totals. This is used in
+			 * reporting of memory context statistics via a sql function. Last
+			 * parameter is not relevant.
+			 */
+			context->methods->stats(context,
+									NULL,
+									NULL,
+									totals, false);
+			break;
+	}
+
+	/* Increment the context count for each of the recursive call */
+	*num_contexts = *num_contexts + 1;
 
 	/*
 	 * Examine children.
@@ -907,7 +942,7 @@ MemoryContextStatsInternal(MemoryContext context, int level,
 			MemoryContextStatsInternal(child, level + 1,
 									   max_level, max_children,
 									   totals,
-									   print_to_stderr);
+									   print_location, num_contexts);
 		}
 	}
 
@@ -926,7 +961,13 @@ MemoryContextStatsInternal(MemoryContext context, int level,
 			child = MemoryContextTraverseNext(child, context);
 		}
 
-		if (print_to_stderr)
+		/*
+		 * Add the count of children contexts which are traversed in the
+		 * non-recursive manner.
+		 */
+		*num_contexts = *num_contexts + ichild;
+
+		if (print_location == PRINT_STATS_TO_STDERR)
 		{
 			for (int i = 0; i < level; i++)
 				fprintf(stderr, "  ");
@@ -939,7 +980,7 @@ MemoryContextStatsInternal(MemoryContext context, int level,
 					local_totals.freechunks,
 					local_totals.totalspace - local_totals.freespace);
 		}
-		else
+		else if (print_location == PRINT_STATS_TO_LOGS)
 			ereport(LOG_SERVER_ONLY,
 					(errhidestmt(true),
 					 errhidecontext(true),
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 1fc19146f46..4bc4474b580 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8597,6 +8597,16 @@
   prorettype => 'bool', proargtypes => 'int4',
   prosrc => 'pg_log_backend_memory_contexts' },
 
+# publishing memory contexts of the specified postgres process
+{ oid => '2173', descr => 'publish memory contexts of the specified backend',
+  proname => 'pg_get_process_memory_contexts', provolatile => 'v',
+  prorows => '100', proretset => 't', proparallel => 'r',
+  prorettype => 'record', proargtypes => 'int4 bool float8',
+  proallargtypes => '{int4,bool,float8,text,text,text,_int4,int4,int8,int8,int8,int8,int8,int4}',
+  proargmodes => '{i,i,i,o,o,o,o,o,o,o,o,o,o,o}',
+  proargnames => '{pid, summary, timeout, name, ident, type, path, level, total_bytes, total_nblocks, free_bytes, free_chunks, used_bytes, num_agg_contexts}',
+  prosrc => 'pg_get_process_memory_contexts' },
+
 # non-persistent series generator
 { oid => '1066', descr => 'non-persistent series generator',
   proname => 'generate_series', prorows => '1000',
diff --git a/src/include/miscadmin.h b/src/include/miscadmin.h
index 1bef98471c3..1e59a7f910f 100644
--- a/src/include/miscadmin.h
+++ b/src/include/miscadmin.h
@@ -96,6 +96,7 @@ extern PGDLLIMPORT volatile sig_atomic_t IdleSessionTimeoutPending;
 extern PGDLLIMPORT volatile sig_atomic_t ProcSignalBarrierPending;
 extern PGDLLIMPORT volatile sig_atomic_t LogMemoryContextPending;
 extern PGDLLIMPORT volatile sig_atomic_t IdleStatsUpdateTimeoutPending;
+extern PGDLLIMPORT volatile sig_atomic_t PublishMemoryContextPending;
 
 extern PGDLLIMPORT volatile sig_atomic_t CheckClientConnectionPending;
 extern PGDLLIMPORT volatile sig_atomic_t ClientConnectionLost;
diff --git a/src/include/storage/lwlock.h b/src/include/storage/lwlock.h
index 08a72569ae5..638407adf39 100644
--- a/src/include/storage/lwlock.h
+++ b/src/include/storage/lwlock.h
@@ -221,6 +221,7 @@ typedef enum BuiltinTrancheIds
 	LWTRANCHE_XACT_SLRU,
 	LWTRANCHE_PARALLEL_VACUUM_DSA,
 	LWTRANCHE_AIO_URING_COMPLETION,
+	LWTRANCHE_MEMORY_CONTEXT_KEYS,
 	LWTRANCHE_FIRST_USER_DEFINED,
 }			BuiltinTrancheIds;
 
diff --git a/src/include/storage/procsignal.h b/src/include/storage/procsignal.h
index afeeb1ca019..345d5a0ecb1 100644
--- a/src/include/storage/procsignal.h
+++ b/src/include/storage/procsignal.h
@@ -35,6 +35,7 @@ typedef enum
 	PROCSIG_WALSND_INIT_STOPPING,	/* ask walsenders to prepare for shutdown  */
 	PROCSIG_BARRIER,			/* global barrier interrupt  */
 	PROCSIG_LOG_MEMORY_CONTEXT, /* ask backend to log the memory contexts */
+	PROCSIG_GET_MEMORY_CONTEXT, /* ask backend to send the memory contexts */
 	PROCSIG_PARALLEL_APPLY_MESSAGE, /* Message from parallel apply workers */
 
 	/* Recovery conflict reasons */
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index 8abc26abce2..01835d56021 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -18,7 +18,10 @@
 #define MEMUTILS_H
 
 #include "nodes/memnodes.h"
-
+#include "storage/condition_variable.h"
+#include "storage/lmgr.h"
+#include "utils/dsa.h"
+#include "lib/dshash.h"
 
 /*
  * MaxAllocSize, MaxAllocHugeSize
@@ -48,6 +51,23 @@
 
 #define AllocHugeSizeIsValid(size)	((Size) (size) <= MaxAllocHugeSize)
 
+/*
+ * Memory Context reporting size limits.
+ */
+
+/* Max length of context name and ident */
+#define MEMORY_CONTEXT_IDENT_SHMEM_SIZE 64
+/* Maximum size (in bytes) of DSA area per process */
+#define MEMORY_CONTEXT_REPORT_MAX_PER_BACKEND  ((size_t) (1 * 1024 * 1024))
+
+/*
+ * Maximum size per context. Actual size may be lower as this assumes the worst
+ * case of deepest path and longest identifiers (name and ident, thus the
+ * multiplication by 2). The path depth is limited to 100 like for memory
+ * context logging.
+ */
+#define MAX_MEMORY_CONTEXT_STATS_SIZE (sizeof(MemoryStatsEntry))
+#define MAX_MEMORY_CONTEXT_STATS_NUM MEMORY_CONTEXT_REPORT_MAX_PER_BACKEND / MAX_MEMORY_CONTEXT_STATS_SIZE
 
 /*
  * Standard top-level memory contexts.
@@ -319,4 +339,74 @@ pg_memory_is_all_zeros(const void *ptr, size_t len)
 	return true;
 }
 
+/* Dynamic shared memory state for statistics per context */
+typedef struct MemoryStatsEntry
+{
+	char		name[MEMORY_CONTEXT_IDENT_SHMEM_SIZE];
+	char		ident[MEMORY_CONTEXT_IDENT_SHMEM_SIZE];
+	int			path[100];
+	NodeTag		type;
+	int			path_length;
+	int			levels;
+	int64		totalspace;
+	int64		nblocks;
+	int64		freespace;
+	int64		freechunks;
+	int			num_agg_stats;
+} MemoryStatsEntry;
+
+/*
+ * Per backend dynamic shared hash entry for memory context statistics
+ * reporting.
+ */
+typedef struct MemoryStatsDSHashEntry
+{
+	char		key[64];
+	ConditionVariable memcxt_cv;
+	int			proc_id;
+	int			total_stats;
+	bool		summary;
+	dsa_pointer memstats_dsa_pointer;
+	TimestampTz stats_timestamp;
+} MemoryStatsDSHashEntry;
+
+static const dshash_parameters memctx_dsh_params = {
+	offsetof(MemoryStatsDSHashEntry, memcxt_cv),
+	sizeof(MemoryStatsDSHashEntry),
+	dshash_strcmp,
+	dshash_strhash,
+	dshash_strcpy
+};
+
+/*
+ * Used for storage of transient identifiers for pg_get_backend_memory_contexts
+ */
+typedef struct MemoryStatsContextId
+{
+	MemoryContext context;
+	int			context_id;
+} MemoryStatsContextId;
+
+/*
+ * This is passed to MemoryContextStatsInternal to determine whether
+ * to print context statistics or not and where to print them logs or
+ * stderr.
+ */
+typedef enum PrintDestination
+{
+	PRINT_STATS_TO_STDERR = 0,
+	PRINT_STATS_TO_LOGS,
+	PRINT_STATS_NONE
+}			PrintDestination;
+
+extern void ProcessGetMemoryContextInterrupt(void);
+extern void HandleGetMemoryContextInterrupt(void);
+extern void MemoryContextKeysShmemInit(void);
+extern Size MemoryContextKeysShmemSize(void);
+extern void MemoryContextStatsInternal(MemoryContext context, int level,
+									   int max_level, int max_children,
+									   MemoryContextCounters *totals,
+									   PrintDestination print_location,
+									   int *num_contexts);
+extern void AtProcExit_memstats_cleanup(int code, Datum arg);
 #endif							/* MEMUTILS_H */
diff --git a/src/test/regress/expected/sysviews.out b/src/test/regress/expected/sysviews.out
index 83228cfca29..ae17d028ed3 100644
--- a/src/test/regress/expected/sysviews.out
+++ b/src/test/regress/expected/sysviews.out
@@ -232,3 +232,22 @@ select * from pg_timezone_abbrevs where abbrev = 'LMT';
  LMT    | @ 7 hours 52 mins 58 secs ago | f
 (1 row)
 
+DO $$
+DECLARE
+    bg_writer_pid int;
+    r RECORD;
+BEGIN
+        SELECT pid from pg_stat_activity where backend_type='background writer'
+	 INTO bg_writer_pid;
+
+        select type, name, ident
+        from pg_get_process_memory_contexts(bg_writer_pid, false, 20)
+	 where path = '{1}' into r;
+	RAISE NOTICE '%', r;
+        select type, name, ident
+        from pg_get_process_memory_contexts(pg_backend_pid(), false, 20)
+	 where path = '{1}' into r;
+	RAISE NOTICE '%', r;
+END $$;
+NOTICE:  (AllocSet,TopMemoryContext,)
+NOTICE:  (AllocSet,TopMemoryContext,)
diff --git a/src/test/regress/sql/sysviews.sql b/src/test/regress/sql/sysviews.sql
index 66179f026b3..d0917b6868e 100644
--- a/src/test/regress/sql/sysviews.sql
+++ b/src/test/regress/sql/sysviews.sql
@@ -101,3 +101,21 @@ select count(distinct utc_offset) >= 24 as ok from pg_timezone_abbrevs;
 -- One specific case we can check without much fear of breakage
 -- is the historical local-mean-time value used for America/Los_Angeles.
 select * from pg_timezone_abbrevs where abbrev = 'LMT';
+
+DO $$
+DECLARE
+    bg_writer_pid int;
+    r RECORD;
+BEGIN
+        SELECT pid from pg_stat_activity where backend_type='background writer'
+	 INTO bg_writer_pid;
+
+        select type, name, ident
+        from pg_get_process_memory_contexts(bg_writer_pid, false, 20)
+	 where path = '{1}' into r;
+	RAISE NOTICE '%', r;
+        select type, name, ident
+        from pg_get_process_memory_contexts(pg_backend_pid(), false, 20)
+	 where path = '{1}' into r;
+	RAISE NOTICE '%', r;
+END $$;
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 83192038571..fedae342032 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -1683,6 +1683,9 @@ MemoryContextData
 MemoryContextId
 MemoryContextMethodID
 MemoryContextMethods
+MemoryStatsContextId
+MemoryStatsEntry
+MemoryStatsDSHashEntry
 MemoryStatsPrintFunc
 MergeAction
 MergeActionState
-- 
2.34.1

