From 0d22a65edac362005d0c810045ae487cafec95c3 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddyforpostgres@gmail.com>
Date: Mon, 15 Nov 2021 02:13:11 +0000
Subject: [PATCH v3] enhance pg_log_backend_memory_contexts to log memory
 contexts of aux procs

Currently pg_log_backend_memory_contexts() doesn't log the memory
contexts of auxiliary processes such as bgwriter, checkpointer, wal
writer, archiver, startup process and wal receiver. It will be useful
to look at the memory contexts of these processes too, for debugging
purposes and better understanding of the memory usage pattern of
these processes. Inside the code, we could use the AuxiliaryPidGetProc()
to get the PGPROC of these processes. Note that, neither
AuxiliaryPidGetProc() nor BackendPidGetProc() can return PGPROC(as
they don't have PGPROC entries at all) entries for the syslogger,
stats collector processes.
---
 doc/src/sgml/func.sgml                       | 14 +++++++---
 src/backend/postmaster/checkpointer.c        |  4 +++
 src/backend/postmaster/interrupt.c           |  5 ++++
 src/backend/postmaster/pgarch.c              |  5 ++++
 src/backend/postmaster/startup.c             |  5 ++++
 src/backend/postmaster/walwriter.c           |  4 +++
 src/backend/utils/adt/mcxtfuncs.c            | 29 ++++++++++++++------
 src/test/regress/expected/misc_functions.out | 11 ++++++++
 src/test/regress/sql/misc_functions.sql      |  9 ++++++
 9 files changed, 74 insertions(+), 12 deletions(-)

diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
index 24447c0017..64091030f7 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -25335,10 +25335,16 @@ SELECT collation for ('foo' COLLATE "de_DE");
         <returnvalue>boolean</returnvalue>
        </para>
        <para>
-        Requests to log the memory contexts of the backend with the
-        specified process ID.  These memory contexts will be logged at
-        <literal>LOG</literal> message level. They will appear in
-        the server log based on the log configuration set
+        Requests to log memory contexts of the <glossterm linkend="glossary-backend">backend</glossterm>
+        or the <glossterm linkend="glossary-wal-sender">WAL sender</glossterm> or
+        the <glossterm linkend="glossary-auxiliary-proc">auxiliary process</glossterm>
+        with the specified process ID. All of the
+        <glossterm linkend="glossary-auxiliary-proc">auxiliary processes</glossterm>
+        are supported except the <glossterm linkend="glossary-logger">logger</glossterm>
+        and the <glossterm linkend="glossary-stats-collector">statistics collector</glossterm>
+        as they are not connected to shared memory the function can not make requests.
+        These memory contexts will be logged at <literal>LOG</literal> message level.
+        They will appear in the server log based on the log configuration set
         (See <xref linkend="runtime-config-logging"/> for more information),
         but will not be sent to the client regardless of
         <xref linkend="guc-client-min-messages"/>.
diff --git a/src/backend/postmaster/checkpointer.c b/src/backend/postmaster/checkpointer.c
index be7366379d..977f3570aa 100644
--- a/src/backend/postmaster/checkpointer.c
+++ b/src/backend/postmaster/checkpointer.c
@@ -581,6 +581,10 @@ HandleCheckpointerInterrupts(void)
 		/* Normal exit from the checkpointer is here */
 		proc_exit(0);			/* done */
 	}
+
+	/* Perform logging of memory contexts of this process */
+	if (LogMemoryContextPending)
+		ProcessLogMemoryContextInterrupt();
 }
 
 /*
diff --git a/src/backend/postmaster/interrupt.c b/src/backend/postmaster/interrupt.c
index dd9136a942..7c1d35a715 100644
--- a/src/backend/postmaster/interrupt.c
+++ b/src/backend/postmaster/interrupt.c
@@ -22,6 +22,7 @@
 #include "storage/latch.h"
 #include "storage/procsignal.h"
 #include "utils/guc.h"
+#include "utils/memutils.h"
 
 volatile sig_atomic_t ConfigReloadPending = false;
 volatile sig_atomic_t ShutdownRequestPending = false;
@@ -43,6 +44,10 @@ HandleMainLoopInterrupts(void)
 
 	if (ShutdownRequestPending)
 		proc_exit(0);
+
+	/* Perform logging of memory contexts of this process */
+	if (LogMemoryContextPending)
+		ProcessLogMemoryContextInterrupt();
 }
 
 /*
diff --git a/src/backend/postmaster/pgarch.c b/src/backend/postmaster/pgarch.c
index 3b33e01d95..eb42682a74 100644
--- a/src/backend/postmaster/pgarch.c
+++ b/src/backend/postmaster/pgarch.c
@@ -50,6 +50,7 @@
 #include "storage/shmem.h"
 #include "storage/spin.h"
 #include "utils/guc.h"
+#include "utils/memutils.h"
 #include "utils/ps_status.h"
 
 
@@ -856,4 +857,8 @@ HandlePgArchInterrupts(void)
 		ConfigReloadPending = false;
 		ProcessConfigFile(PGC_SIGHUP);
 	}
+
+	/* Perform logging of memory contexts of this process */
+	if (LogMemoryContextPending)
+		ProcessLogMemoryContextInterrupt();
 }
diff --git a/src/backend/postmaster/startup.c b/src/backend/postmaster/startup.c
index 47ec737888..2b35a82903 100644
--- a/src/backend/postmaster/startup.c
+++ b/src/backend/postmaster/startup.c
@@ -32,6 +32,7 @@
 #include "storage/procsignal.h"
 #include "storage/standby.h"
 #include "utils/guc.h"
+#include "utils/memutils.h"
 #include "utils/timeout.h"
 
 
@@ -200,6 +201,10 @@ HandleStartupProcInterrupts(void)
 	/* Process barrier events */
 	if (ProcSignalBarrierPending)
 		ProcessProcSignalBarrier();
+
+	/* Perform logging of memory contexts of this process */
+	if (LogMemoryContextPending)
+		ProcessLogMemoryContextInterrupt();
 }
 
 
diff --git a/src/backend/postmaster/walwriter.c b/src/backend/postmaster/walwriter.c
index 626fae8454..19c9118fd8 100644
--- a/src/backend/postmaster/walwriter.c
+++ b/src/backend/postmaster/walwriter.c
@@ -306,4 +306,8 @@ HandleWalWriterInterrupts(void)
 
 		proc_exit(0);
 	}
+
+	/* Perform logging of memory contexts of this process */
+	if (LogMemoryContextPending)
+		ProcessLogMemoryContextInterrupt();
 }
diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
index 6ddbf70b30..67598ddbb3 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -176,17 +176,26 @@ pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
 {
 	int			pid = PG_GETARG_INT32(0);
 	PGPROC	   *proc;
+	bool		is_aux_proc = false;
+	BackendId   backendId = InvalidBackendId;
 
 	proc = BackendPidGetProc(pid);
 
+	/* See if the process with given pid is an auxiliary process. */
+	if (proc == NULL)
+	{
+		proc = AuxiliaryPidGetProc(pid);
+		is_aux_proc = true;
+	}
+
 	/*
-	 * BackendPidGetProc returns 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 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.
+	 * BackendPidGetProc or AuxiliaryPidGetProc returns 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 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.
 	 */
 	if (proc == NULL)
 	{
@@ -199,7 +208,11 @@ pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
 		PG_RETURN_BOOL(false);
 	}
 
-	if (SendProcSignal(pid, PROCSIG_LOG_MEMORY_CONTEXT, proc->backendId) < 0)
+	/* Only regular backends will have valid backend id, auxiliary processes don't. */
+	if (!is_aux_proc)
+		backendId = proc->backendId;
+
+	if (SendProcSignal(pid, PROCSIG_LOG_MEMORY_CONTEXT, backendId) < 0)
 	{
 		/* Again, just a warning to allow loops */
 		ereport(WARNING,
diff --git a/src/test/regress/expected/misc_functions.out b/src/test/regress/expected/misc_functions.out
index 71d316cad3..ca52779f85 100644
--- a/src/test/regress/expected/misc_functions.out
+++ b/src/test/regress/expected/misc_functions.out
@@ -147,6 +147,17 @@ SELECT pg_log_backend_memory_contexts(pg_backend_pid());
  t
 (1 row)
 
+CREATE FUNCTION memcxt_get_proc_pid(text)
+RETURNS int
+LANGUAGE SQL
+AS 'SELECT pid FROM pg_stat_activity WHERE backend_type = $1';
+SELECT pg_log_backend_memory_contexts(memcxt_get_proc_pid('checkpointer'));
+ pg_log_backend_memory_contexts 
+--------------------------------
+ t
+(1 row)
+
+DROP FUNCTION memcxt_get_proc_pid(text);
 CREATE ROLE regress_log_memory;
 SELECT has_function_privilege('regress_log_memory',
   'pg_log_backend_memory_contexts(integer)', 'EXECUTE'); -- no
diff --git a/src/test/regress/sql/misc_functions.sql b/src/test/regress/sql/misc_functions.sql
index 8c23874b3f..a9948e59bb 100644
--- a/src/test/regress/sql/misc_functions.sql
+++ b/src/test/regress/sql/misc_functions.sql
@@ -41,6 +41,15 @@ SELECT num_nulls();
 
 SELECT pg_log_backend_memory_contexts(pg_backend_pid());
 
+CREATE FUNCTION memcxt_get_proc_pid(text)
+RETURNS int
+LANGUAGE SQL
+AS 'SELECT pid FROM pg_stat_activity WHERE backend_type = $1';
+
+SELECT pg_log_backend_memory_contexts(memcxt_get_proc_pid('checkpointer'));
+
+DROP FUNCTION memcxt_get_proc_pid(text);
+
 CREATE ROLE regress_log_memory;
 
 SELECT has_function_privilege('regress_log_memory',
-- 
2.25.1

