From a98f27ccecd62867ed7aeb9be166488fffb4304d Mon Sep 17 00:00:00 2001
From: Bertrand Drouvot <bertranddrouvot.pg@gmail.com>
Date: Tue, 25 Feb 2025 09:03:55 +0000
Subject: [PATCH v11 1/2] Add the pg_stat_get_backend_stats() helper for
 pg_stat_get_backend_io()

This commit adds pg_stat_get_backend_stats(), a helper routine for
pg_stat_get_backend_io(), that returns the backend stats based on a pid passed
as an argument.

This will be used in a follow-up commit that uses the same logic to return the
per backend WAL stats.
---
 src/backend/utils/activity/pgstat_backend.c | 52 +++++++++++++++++++++
 src/backend/utils/adt/pgstatfuncs.c         | 30 +-----------
 src/include/pgstat.h                        |  1 +
 3 files changed, 54 insertions(+), 29 deletions(-)
  58.1% src/backend/utils/activity/
  38.0% src/backend/utils/adt/
   3.7% src/include/

diff --git a/src/backend/utils/activity/pgstat_backend.c b/src/backend/utils/activity/pgstat_backend.c
index 338da73a9a9..6f17e1ff39f 100644
--- a/src/backend/utils/activity/pgstat_backend.c
+++ b/src/backend/utils/activity/pgstat_backend.c
@@ -26,6 +26,8 @@
 
 #include "access/xlog.h"
 #include "storage/bufmgr.h"
+#include "storage/proc.h"
+#include "storage/procarray.h"
 #include "utils/memutils.h"
 #include "utils/pgstat_internal.h"
 
@@ -82,6 +84,56 @@ pgstat_fetch_stat_backend(ProcNumber procNumber)
 	return backend_entry;
 }
 
+/*
+ * Returns statistics of a backend by pid.
+ *
+ * It adds extra checks as compared to pgstat_fetch_stat_backend() to ensure
+ * that the backend is not gone. Also, if not NULL, bktype is populated as
+ * pg_stat_get_backend_io() needs it.
+ */
+PgStat_Backend *
+pg_stat_get_backend_stats(int pid, BackendType *bktype)
+{
+
+	PGPROC	   *proc;
+	PgBackendStatus *beentry;
+	ProcNumber	procNumber;
+	PgStat_Backend *backend_stats;
+
+	proc = BackendPidGetProc(pid);
+
+	/*
+	 * This could be an auxiliary process but these do not report backend
+	 * statistics due to pgstat_tracks_backend_bktype(), so there is no need
+	 * for an extra call to AuxiliaryPidGetProc().
+	 */
+	if (!proc)
+		return NULL;
+
+	procNumber = GetNumberFromPGProc(proc);
+
+	beentry = pgstat_get_beentry_by_proc_number(procNumber);
+	if (!beentry)
+		return NULL;
+
+	backend_stats = pgstat_fetch_stat_backend(procNumber);
+	if (!backend_stats)
+		return NULL;
+
+	/* if PID does not match, leave */
+	if (beentry->st_procpid != pid)
+		return NULL;
+
+	/* backend may be gone, so recheck in case */
+	if (beentry->st_backendType == B_INVALID)
+		return NULL;
+
+	if (bktype)
+		*bktype = beentry->st_backendType;
+
+	return backend_stats;
+}
+
 /*
  * Flush out locally pending backend IO statistics.  Locking is managed
  * by the caller.
diff --git a/src/backend/utils/adt/pgstatfuncs.c b/src/backend/utils/adt/pgstatfuncs.c
index efb6d0032af..ea91c8fc9d5 100644
--- a/src/backend/utils/adt/pgstatfuncs.c
+++ b/src/backend/utils/adt/pgstatfuncs.c
@@ -1576,46 +1576,18 @@ pg_stat_get_backend_io(PG_FUNCTION_ARGS)
 	ReturnSetInfo *rsinfo;
 	BackendType bktype;
 	int			pid;
-	PGPROC	   *proc;
-	ProcNumber	procNumber;
 	PgStat_Backend *backend_stats;
 	PgStat_BktypeIO *bktype_stats;
-	PgBackendStatus *beentry;
 
 	InitMaterializedSRF(fcinfo, 0);
 	rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
 
 	pid = PG_GETARG_INT32(0);
-	proc = BackendPidGetProc(pid);
-
-	/*
-	 * This could be an auxiliary process but these do not report backend
-	 * statistics due to pgstat_tracks_backend_bktype(), so there is no need
-	 * for an extra call to AuxiliaryPidGetProc().
-	 */
-	if (!proc)
-		return (Datum) 0;
-
-	procNumber = GetNumberFromPGProc(proc);
+	backend_stats = pg_stat_get_backend_stats(pid, &bktype);
 
-	beentry = pgstat_get_beentry_by_proc_number(procNumber);
-	if (!beentry)
-		return (Datum) 0;
-
-	backend_stats = pgstat_fetch_stat_backend(procNumber);
 	if (!backend_stats)
 		return (Datum) 0;
 
-	bktype = beentry->st_backendType;
-
-	/* if PID does not match, leave */
-	if (beentry->st_procpid != pid)
-		return (Datum) 0;
-
-	/* backend may be gone, so recheck in case */
-	if (bktype == B_INVALID)
-		return (Datum) 0;
-
 	bktype_stats = &backend_stats->io_stats;
 
 	/*
diff --git a/src/include/pgstat.h b/src/include/pgstat.h
index 67656264b62..e8e7d95b334 100644
--- a/src/include/pgstat.h
+++ b/src/include/pgstat.h
@@ -554,6 +554,7 @@ extern void pgstat_count_backend_io_op(IOObject io_object,
 									   IOOp io_op, uint32 cnt,
 									   uint64 bytes);
 extern PgStat_Backend *pgstat_fetch_stat_backend(ProcNumber procNumber);
+extern PgStat_Backend *pg_stat_get_backend_stats(int pid, BackendType *bktype);
 extern bool pgstat_tracks_backend_bktype(BackendType bktype);
 extern void pgstat_create_backend(ProcNumber procnum);
 
-- 
2.34.1