From 48d715a4a8b667628bb8ccf573e0003eb8f54633 Mon Sep 17 00:00:00 2001
From: Bertrand Drouvot <bertranddrouvot.pg@gmail.com>
Date: Fri, 17 Jan 2025 06:44:47 +0000
Subject: [PATCH v1] Rework per backend pending stats

9aea73fc61d added per backend pending statistics but not all of them are well
suited to memory allocation and have to be outside of critical sections.

For those (the only current one is I/O statistics but WAL statistics is under
discussion), let's rely on a new PendingBackendStats instead.
---
 src/backend/utils/activity/pgstat.c         | 27 ++++++++++++++++++-
 src/backend/utils/activity/pgstat_backend.c | 29 +++++++++++++++------
 src/backend/utils/activity/pgstat_io.c      | 14 +++-------
 src/include/pgstat.h                        | 12 +++++++++
 4 files changed, 62 insertions(+), 20 deletions(-)
  89.0% src/backend/utils/activity/
  10.9% src/include/

diff --git a/src/backend/utils/activity/pgstat.c b/src/backend/utils/activity/pgstat.c
index 34520535d54..27aa0491001 100644
--- a/src/backend/utils/activity/pgstat.c
+++ b/src/backend/utils/activity/pgstat.c
@@ -1369,6 +1369,7 @@ static bool
 pgstat_flush_pending_entries(bool nowait)
 {
 	bool		have_pending = false;
+	bool		backend_have_pending = false;
 	dlist_node *cur = NULL;
 
 	/*
@@ -1416,9 +1417,33 @@ pgstat_flush_pending_entries(bool nowait)
 		cur = next;
 	}
 
+	/*
+	 * There is a special case for some pending stats that are tracked in
+	 * PendingBackendStats. It's possible that those have not been flushed
+	 * above, hence the extra check here.
+	 */
+	if (!pg_memory_is_all_zeros(&PendingBackendStats,
+								sizeof(struct PgStat_BackendPending)))
+	{
+		PgStat_EntryRef *entry_ref;
+
+		entry_ref = pgstat_get_entry_ref(PGSTAT_KIND_BACKEND, InvalidOid,
+										 MyProcNumber, false, NULL);
+
+		if (entry_ref)
+		{
+			PgStat_HashKey key = entry_ref->shared_entry->key;
+			PgStat_Kind kind = key.kind;
+			const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
+
+			if (!kind_info->flush_pending_cb(entry_ref, nowait))
+				backend_have_pending = true;
+		}
+	}
+
 	Assert(dlist_is_empty(&pgStatPending) == !have_pending);
 
-	return have_pending;
+	return (have_pending || backend_have_pending);
 }
 
 
diff --git a/src/backend/utils/activity/pgstat_backend.c b/src/backend/utils/activity/pgstat_backend.c
index 79e4d0a3053..ca5cec348b2 100644
--- a/src/backend/utils/activity/pgstat_backend.c
+++ b/src/backend/utils/activity/pgstat_backend.c
@@ -22,8 +22,11 @@
 
 #include "postgres.h"
 
+#include "utils/memutils.h"
 #include "utils/pgstat_internal.h"
 
+PgStat_BackendPending PendingBackendStats = {0};
+
 /*
  * Returns statistics of a backend by proc number.
  */
@@ -46,14 +49,20 @@ static void
 pgstat_flush_backend_entry_io(PgStat_EntryRef *entry_ref)
 {
 	PgStatShared_Backend *shbackendent;
-	PgStat_BackendPending *pendingent;
 	PgStat_BktypeIO *bktype_shstats;
-	PgStat_PendingIO *pending_io;
+	PgStat_PendingIO pending_io;
+
+	/*
+	 * This function can be called even if nothing at all has happened. In
+	 * this case, avoid unnecessarily modifying the stats entry.
+	 */
+	if (pg_memory_is_all_zeros(&PendingBackendStats.pending_io,
+							   sizeof(struct PgStat_PendingIO)))
+		return;
 
 	shbackendent = (PgStatShared_Backend *) entry_ref->shared_stats;
-	pendingent = (PgStat_BackendPending *) entry_ref->pending;
 	bktype_shstats = &shbackendent->stats.io_stats;
-	pending_io = &pendingent->pending_io;
+	pending_io = PendingBackendStats.pending_io;
 
 	for (int io_object = 0; io_object < IOOBJECT_NUM_TYPES; io_object++)
 	{
@@ -64,17 +73,21 @@ pgstat_flush_backend_entry_io(PgStat_EntryRef *entry_ref)
 				instr_time	time;
 
 				bktype_shstats->counts[io_object][io_context][io_op] +=
-					pending_io->counts[io_object][io_context][io_op];
+					pending_io.counts[io_object][io_context][io_op];
 				bktype_shstats->bytes[io_object][io_context][io_op] +=
-					pending_io->bytes[io_object][io_context][io_op];
-
-				time = pending_io->pending_times[io_object][io_context][io_op];
+					pending_io.bytes[io_object][io_context][io_op];
+				time = pending_io.pending_times[io_object][io_context][io_op];
 
 				bktype_shstats->times[io_object][io_context][io_op] +=
 					INSTR_TIME_GET_MICROSEC(time);
 			}
 		}
 	}
+
+	/*
+	 * Clear out the statistics buffer, so it can be re-used.
+	 */
+	MemSet(&PendingBackendStats.pending_io, 0, sizeof(PgStat_PendingIO));
 }
 
 /*
diff --git a/src/backend/utils/activity/pgstat_io.c b/src/backend/utils/activity/pgstat_io.c
index 027aad8b24e..b8f19515ae1 100644
--- a/src/backend/utils/activity/pgstat_io.c
+++ b/src/backend/utils/activity/pgstat_io.c
@@ -75,11 +75,8 @@ pgstat_count_io_op(IOObject io_object, IOContext io_context, IOOp io_op,
 
 	if (pgstat_tracks_backend_bktype(MyBackendType))
 	{
-		PgStat_BackendPending *entry_ref;
-
-		entry_ref = pgstat_prep_backend_pending(MyProcNumber);
-		entry_ref->pending_io.counts[io_object][io_context][io_op] += cnt;
-		entry_ref->pending_io.bytes[io_object][io_context][io_op] += bytes;
+		PendingBackendStats.pending_io.counts[io_object][io_context][io_op] += cnt;
+		PendingBackendStats.pending_io.bytes[io_object][io_context][io_op] += bytes;
 	}
 
 	PendingIOStats.counts[io_object][io_context][io_op] += cnt;
@@ -146,13 +143,8 @@ pgstat_count_io_op_time(IOObject io_object, IOContext io_context, IOOp io_op,
 					   io_time);
 
 		if (pgstat_tracks_backend_bktype(MyBackendType))
-		{
-			PgStat_BackendPending *entry_ref;
-
-			entry_ref = pgstat_prep_backend_pending(MyProcNumber);
-			INSTR_TIME_ADD(entry_ref->pending_io.pending_times[io_object][io_context][io_op],
+			INSTR_TIME_ADD(PendingBackendStats.pending_io.pending_times[io_object][io_context][io_op],
 						   io_time);
-		}
 	}
 
 	pgstat_count_io_op(io_object, io_context, io_op, cnt, bytes);
diff --git a/src/include/pgstat.h b/src/include/pgstat.h
index 2d40fe3e70f..3eb7b6614fe 100644
--- a/src/include/pgstat.h
+++ b/src/include/pgstat.h
@@ -826,5 +826,17 @@ extern PGDLLIMPORT SessionEndType pgStatSessionEndCause;
 /* updated directly by backends and background processes */
 extern PGDLLIMPORT PgStat_PendingWalStats PendingWalStats;
 
+/*
+ * Variables in pgstat_backend.c
+ */
+
+/* updated directly by backends */
+
+/*
+ * Some pending statistics are not well suited to memory allocation and have
+ * to be outside of critical sections. For those, let's rely on this variable
+ * instead.
+ */
+extern PGDLLIMPORT PgStat_BackendPending PendingBackendStats;
 
 #endif							/* PGSTAT_H */
-- 
2.34.1

