From 90ffa6d32a9cd0034a2fd97c31991199f2089698 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 v2 1/2] 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/include/pgstat.h                        | 12 +++++++++
 src/backend/utils/activity/pgstat.c         | 27 ++++++++++++++++++-
 src/backend/utils/activity/pgstat_backend.c | 29 +++++++++++++++------
 src/backend/utils/activity/pgstat_io.c      | 14 +++-------
 4 files changed, 62 insertions(+), 20 deletions(-)

diff --git a/src/include/pgstat.h b/src/include/pgstat.h
index 2d40fe3e70..3eb7b6614f 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 */
diff --git a/src/backend/utils/activity/pgstat.c b/src/backend/utils/activity/pgstat.c
index 34520535d5..27aa049100 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 79e4d0a305..ca5cec348b 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 027aad8b24..b8f19515ae 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);
-- 
2.47.1

