From d9e8a1488472db17bcb3b275488232e5db94d018 Mon Sep 17 00:00:00 2001
From: Daniel Gustafsson <dgustafsson@postgresql.org>
Date: Tue, 25 Mar 2025 13:47:52 +0100
Subject: [PATCH v19 3/3] Review comments and fixups

---
 doc/src/sgml/func.sgml            |  62 +++++++++++-------
 src/backend/utils/adt/mcxtfuncs.c | 103 +++++++++++++++---------------
 src/backend/utils/mmgr/mcxt.c     |  78 +++++++++++-----------
 3 files changed, 131 insertions(+), 112 deletions(-)

diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
index 1ab2ce12662..f8a3b9d0cf6 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -28580,34 +28580,36 @@ acl      | {postgres=arwdDxtm/postgres,foo=r/postgres}
        </para>
        <para>
         This function handles requests to display the memory contexts of a
-        PostgreSQL process with the specified process ID (PID). It takes three
-        arguments: <parameter>PID</parameter>, <parameter>get_summary</parameter>
-        and <parameter>num_of_tries</parameter>. The function can send requests
+        <productname>PostgreSQL</productname> process with the specified process ID.  It takes three
+        arguments: <parameter>pid</parameter>, <parameter>get_summary</parameter>
+        and <parameter>num_of_tries</parameter>.  The function can send requests
         to both backend and auxiliary processes.
-
+       </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
+        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 num_agg_contexts
+        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. The num_agg_contexts value is typically 1, meaning that each
+        statistics.  The <literal>num_agg_contexts</literal> value is typically 1, meaning that each
         context's statistics are displayed separately.
-
-        When <parameter>get_summary</parameter> is set to true, statistics
+       </para>
+       <para>
+        When <parameter>get_summary</parameter> is set to <literal>true</literal>, statistics
         for memory contexts at levels 1 and 2 are displayed, with level 1
         representing the root node (i.e., TopMemoryContext).
         Each level 2 context's statistics represent an aggregate of all its
-        child contexts' statistics, with num_agg_contexts indicating the number
+        child contexts' statistics, with <literal>num_agg_contexts</literal> indicating the number
         of these aggregated child contexts.
-
-        When <parameter>get_summary</parameter> is set to false, the
+        When <parameter>get_summary</parameter> is set to <literal>false</literal>, the
         num_agg_contexts value is 1, indicating that individual statistics are
         being displayed.
-
+       </para>
+       <para>
         <parameter>num_of_tries</parameter> indicates the number of times
-        the client will wait for the latest statistics. The wait per try is 1
-        second. This parameter can be increased if the user anticipates a delay
-        in the response from the reporting process. Conversely, if users are
+        the client will wait for the latest statistics.  The wait per try is 1
+        second.  This parameter can be increased if the user anticipates a delay
+        in the response from the reporting process.  Conversely, if users are
         frequently and periodically querying the process for statistics, or if
         there are concurrent requests for statistics of the same process,
         lowering the parameter might help achieve a faster response.
@@ -28740,13 +28742,29 @@ postgres=# SELECT * FROM pg_get_process_memory_contexts(
   (SELECT pid FROM pg_stat_activity
     WHERE backend_type = 'checkpointer')
   , false, 5) LIMIT 1;
-       name       | ident |   type   | path | total_bytes | total_nblocks | free_bytes | free_chunks | used_bytes | num_
-agg_contexts |         stats_timestamp
-------------------+-------+----------+------+-------------+---------------+------------+-------------+------------+-----
--------------+----------------------------------
- TopMemoryContext |       | AllocSet | {1}  |      102664 |             4 |       3008 |           2 |      99656 |
-           1 | 2025-03-04 10:01:57.590543+05:30
+-[ 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
+stats_timestamp  | 2025-03-24 13:55:47.796698+01
 </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>
diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
index 462c4e48cf0..f9145b9faaa 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -330,7 +330,7 @@ pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
  * 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 that can fit in the allocated
+ * 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.
  *
@@ -354,6 +354,7 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 	bool		get_summary = PG_GETARG_BOOL(1);
 	PGPROC	   *proc;
 	ProcNumber	procNumber = INVALID_PROC_NUMBER;
+	bool		proc_is_aux = false;
 	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
 	dsa_area   *area;
 	MemoryContextEntry *memctx_info;
@@ -368,25 +369,25 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 	if (!has_privs_of_role(GetUserId(), ROLE_PG_READ_ALL_STATS))
 		ereport(ERROR,
 				errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
-				errmsg("memory context statistics privilege error"));
+				errmsg("permission denied to extract memory context statistics"));
 
 	InitMaterializedSRF(fcinfo, 0);
 
 	/*
-	 * See if the process with given pid is a backend or an auxiliary process.
+	 * 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; but by the time we reach kill(), a process for which we
-	 * get a valid proc here might have terminated on its own.  There's no way
-	 * to acquire a lock on an arbitrary process to prevent that. But since
-	 * this mechanism is usually used to debug a backend or an auxiliary
-	 * process running and consuming lots of memory, that it might end on its
-	 * own first and its memory contexts are not logged is not a problem.
+	 * 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)
 	{
@@ -395,8 +396,7 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 		 * if one backend terminated on its own during the run.
 		 */
 		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();
 	}
 
@@ -409,8 +409,8 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 	curr_timestamp = GetCurrentTimestamp();
 
 	/*
-	 * Send a signal to a postgresql process, informing it we want it to
-	 * produce information about memory contexts.
+	 * 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)
 	{
@@ -442,9 +442,6 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 		 * process to finish publishing statistics.
 		 */
 		LWLockAcquire(&memCtxState[procNumber].lw_lock, LW_EXCLUSIVE);
-		msecs =
-			TimestampDifferenceMilliseconds(curr_timestamp,
-											memCtxState[procNumber].stats_timestamp);
 
 		/*
 		 * Note in procnumber.h file says that a procNumber can be re-used for
@@ -460,29 +457,38 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 			 * statistics timestamp being newer than the current request
 			 * timestamp.
 			 */
+			msecs = TimestampDifferenceMilliseconds(curr_timestamp,
+													memCtxState[procNumber].stats_timestamp);
+
 			if (DsaPointerIsValid(memCtxState[procNumber].memstats_dsa_pointer)
 				&& msecs > 0)
 				break;
-
 		}
 		LWLockRelease(&memCtxState[procNumber].lw_lock);
 
 		/*
 		 * Recheck the state of the backend before sleeping on the condition
-		 * variable
+		 * variable to ensure the process is still alive.  Only check the
+		 * relevant process type based on the earlier PID check.
 		 */
-		proc = BackendPidGetProc(pid);
-
-#define MEMSTATS_WAIT_TIMEOUT 1000
-		if (proc == NULL)
+		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 not a PostgreSQL server process",
-						   pid));
+					errmsg("PID %d is not a PostgreSQL server process", pid));
 			PG_RETURN_NULL();
 		}
+
+#define MEMSTATS_WAIT_TIMEOUT 1000
+
 		if (ConditionVariableTimedSleep(&memCtxState[procNumber].memctx_cv,
 										MEMSTATS_WAIT_TIMEOUT,
 										WAIT_EVENT_MEM_CTX_PUBLISH))
@@ -503,12 +509,11 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 					PG_RETURN_NULL();
 				}
 			}
-			ereport(LOG,
-					errmsg("Wait for %d process to publish stats timed out, trying again",
+			ereport(DEBUG1,
+					errmsg("timed out waiting for process with PID %d to publish stats, retrying",
 						   pid));
 			num_retries = num_retries + 1;
 		}
-
 	}
 	/* We should land here only with a valid DSA handle */
 	Assert(memCtxArea->memstats_dsa_handle != DSA_HANDLE_INVALID);
@@ -519,8 +524,8 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 	 *
 	 * Read statistics of top level 1 and 2 contexts, if get_summary is true.
 	 */
-	memctx_info = (MemoryContextEntry *) dsa_get_address(area,
-														 memCtxState[procNumber].memstats_dsa_pointer);
+	memctx_info = (MemoryContextEntry *)
+		dsa_get_address(area, memCtxState[procNumber].memstats_dsa_pointer);
 
 #define PG_GET_PROCESS_MEMORY_CONTEXTS_COLS	12
 	for (int i = 0; i < memCtxState[procNumber].total_stats; i++)
@@ -543,6 +548,7 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 		}
 		else
 			nulls[0] = true;
+
 		if (DsaPointerIsValid(memctx_info[i].ident))
 		{
 			ident = (char *) dsa_get_address(area, memctx_info[i].ident);
@@ -563,11 +569,11 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 			path_datum_array = (Datum *) dsa_get_address(area, memctx_info[i].path);
 			path_array = construct_array_builtin(path_datum_array,
 												 path_length, INT4OID);
-
 			values[3] = PointerGetDatum(path_array);
 		}
 		else
 			nulls[3] = true;
+
 		values[4] = Int32GetDatum(path_length); /* level */
 		values[5] = Int64GetDatum(memctx_info[i].totalspace);
 		values[6] = Int64GetDatum(memctx_info[i].nblocks);
@@ -589,18 +595,6 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 	PG_RETURN_NULL();
 }
 
-/*
- * Shared memory sizing for reporting memory context information.
- */
-static Size
-MemCtxShmemSize(void)
-{
-	Size		TotalProcs =
-		add_size(MaxBackends, add_size(NUM_AUXILIARY_PROCS, max_prepared_xacts));
-
-	return mul_size(TotalProcs, sizeof(MemoryContextBackendState));
-}
-
 /*
  * Init shared memory for reporting memory context information.
  */
@@ -608,12 +602,16 @@ void
 MemCtxBackendShmemInit(void)
 {
 	bool		found;
-	Size		TotalProcs =
-		add_size(MaxBackends, add_size(NUM_AUXILIARY_PROCS, max_prepared_xacts));
+	Size		TotalProcs;
+
+	TotalProcs = add_size(MaxBackends, NUM_AUXILIARY_PROCS);
+	TotalProcs = add_size(TotalProcs, max_prepared_xacts);
+
+	memCtxState = (MemoryContextBackendState *)
+		ShmemInitStruct("MemoryContextBackendState",
+						mul_size(TotalProcs, sizeof(MemoryContextBackendState)),
+						&found);
 
-	memCtxState = (MemoryContextBackendState *) ShmemInitStruct("MemoryContextBackendState",
-																MemCtxShmemSize(),
-																&found);
 	if (!IsUnderPostmaster)
 	{
 		Assert(!found);
@@ -622,8 +620,7 @@ MemCtxBackendShmemInit(void)
 		{
 			ConditionVariableInit(&memCtxState[i].memctx_cv);
 
-			LWLockInitialize(&memCtxState[i].lw_lock,
-							 LWLockNewTrancheId());
+			LWLockInitialize(&memCtxState[i].lw_lock, LWLockNewTrancheId());
 			LWLockRegisterTranche(memCtxState[i].lw_lock.tranche,
 								  "mem_context_backend_stats_reporting");
 
@@ -637,16 +634,16 @@ MemCtxBackendShmemInit(void)
 }
 
 /*
- * Initialize shared memory for displaying memory
- * context statistics
+ * Initialize shared memory for displaying memory context statistics
  */
 void
 MemCtxShmemInit(void)
 {
 	bool		found;
 
-	memCtxArea = (MemoryContextState *) ShmemInitStruct("MemoryContextState", sizeof(MemoryContextState),
-														&found);
+	memCtxArea = (MemoryContextState *)
+		ShmemInitStruct("MemoryContextState", sizeof(MemoryContextState), &found);
+
 	if (!IsUnderPostmaster)
 	{
 		Assert(!found);
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 214330aa7a5..4db48230241 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -19,7 +19,6 @@
  *-------------------------------------------------------------------------
  */
 
-#include <math.h>
 #include "postgres.h"
 
 #include "mb/pg_wchar.h"
@@ -922,38 +921,39 @@ MemoryContextStatsInternal(MemoryContext context, int level,
 	check_stack_depth();
 	Assert(MemoryContextIsValid(context));
 
-	if (print_location == PRINT_STATS_TO_STDERR)
-	{
-		/* Examine the context itself */
-		context->methods->stats(context,
-								MemoryContextStatsPrint,
-								&level,
-								totals, true);
-	}
-	else if (print_location == PRINT_STATS_TO_LOGS)
+	switch (print_location)
 	{
+		case PRINT_STATS_TO_STDERR:
+			/* Examine the context itself */
+			context->methods->stats(context,
+									MemoryContextStatsPrint,
+									&level,
+									totals, true);
+			break;
 
-		/* Examine the context itself */
-		context->methods->stats(context,
-								MemoryContextStatsPrint,
-								&level,
-								totals, false);
-	}
+		case PRINT_STATS_TO_LOGS:
+			/* Examine the context itself */
+			context->methods->stats(context,
+									MemoryContextStatsPrint,
+									&level,
+									totals, false);
+			break;
 
-	/*
-	 * Do not print the statistics if print_to_stderr 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.
-	 */
-	else
-	{
-		Assert(print_location == PRINT_STATS_NONE);
-		/* Examine the context itself */
-		context->methods->stats(context,
-								NULL,
-								NULL,
-								totals, false);
+		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;
 
@@ -1601,6 +1601,7 @@ ProcessGetMemoryContextInterrupt(void)
 		memCtxState[idx].total_stats = ctx_id;
 		goto cleanup;
 	}
+
 	foreach_ptr(MemoryContextData, cur, contexts)
 	{
 		List	   *path = NIL;
@@ -1665,6 +1666,7 @@ ProcessGetMemoryContextInterrupt(void)
 		 */
 		memCtxState[idx].total_stats = num_individual_stats + 1;
 	}
+
 cleanup:
 
 	/*
@@ -1680,8 +1682,8 @@ cleanup:
 }
 
 /*
- * Append the transient context_id of this context and each of
- * its ancestors to a list, in order to compute a 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)
@@ -1763,6 +1765,8 @@ PublishMemoryContext(MemoryContextEntry *memctx_info, int curr_id,
 	char	   *ident;
 	Datum	   *path_array;
 
+	Assert(MemoryContextIsValid(context));
+
 	if (context->name != NULL)
 	{
 		Assert(strlen(context->name) < MEMORY_CONTEXT_IDENT_SHMEM_SIZE);
@@ -1817,6 +1821,7 @@ PublishMemoryContext(MemoryContextEntry *memctx_info, int curr_id,
 	}
 	else
 		memctx_info[curr_id].ident = InvalidDsaPointer;
+
 	/* Allocate dsa memory for storing path information */
 	if (path == NIL)
 		memctx_info[curr_id].path = InvalidDsaPointer;
@@ -1869,11 +1874,11 @@ AtProcExit_memstats_dsa_free(int code, Datum arg)
 	dsm_segment *dsm_seg = NULL;
 	dsa_area   *area = NULL;
 
-	if (memCtxArea->memstats_dsa_handle != DSA_HANDLE_INVALID)
-		dsm_seg = dsm_find_mapping(memCtxArea->memstats_dsa_handle);
-	else
+	if (memCtxArea->memstats_dsa_handle == DSA_HANDLE_INVALID)
 		return;
 
+	dsm_seg = dsm_find_mapping(memCtxArea->memstats_dsa_handle);
+
 	LWLockAcquire(&memCtxState[idx].lw_lock, LW_EXCLUSIVE);
 
 	if (!DsaPointerIsValid(memCtxState[idx].memstats_dsa_pointer))
@@ -1883,10 +1888,9 @@ AtProcExit_memstats_dsa_free(int code, Datum arg)
 	}
 
 	/* If the dsm mapping could not be found, attach to the area */
-	if (dsm_seg == NULL)
-		area = dsa_attach(memCtxArea->memstats_dsa_handle);
-	else
+	if (dsm_seg != NULL)
 		return;
+	area = dsa_attach(memCtxArea->memstats_dsa_handle);
 
 	/*
 	 * Free the memory context statistics, free the name, ident and path
-- 
2.39.3 (Apple Git-146)

