From 0c924a72f586c385f6ab1a22174b9c3b1cf2dd08 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Fri, 12 Sep 2025 16:09:51 +0900
Subject: [PATCH 1/2] Add support for entry counting in pgstats

Stats kinds can set an option call track_counts, that will make pgstats
track the number of entries that exist in the shared hashtable.

As there is only one code path where a new entry is added, and one code
path where entries are removed, the count tracking is rather
straight-forward.  Reads of these counters are optimistic, and may
change across two calls.

A use case of this facility is PGSS, where we need to be able to cap the
number of entries that would be stored in the shared hashtable, based on
its "max" GUC.
---
 src/include/utils/pgstat_internal.h       | 26 +++++++++++++
 src/backend/utils/activity/pgstat.c       |  6 +++
 src/backend/utils/activity/pgstat_shmem.c | 47 +++++++++++++++++------
 3 files changed, 67 insertions(+), 12 deletions(-)

diff --git a/src/include/utils/pgstat_internal.h b/src/include/utils/pgstat_internal.h
index 6cf00008f633..8762a06081f8 100644
--- a/src/include/utils/pgstat_internal.h
+++ b/src/include/utils/pgstat_internal.h
@@ -216,6 +216,12 @@ typedef struct PgStat_KindInfo
 	/* Should stats be written to the on-disk stats file? */
 	bool		write_to_file:1;
 
+	/*
+	 * Should the number of entries be tracked?  For variable-numbered stats,
+	 * to update its PgStat_ShmemControl.entry_counts.
+	 */
+	bool		track_counts:1;
+
 	/*
 	 * The size of an entry in the shared stats hash table (pointed to by
 	 * PgStatShared_HashEntry->body).  For fixed-numbered statistics, this is
@@ -485,6 +491,15 @@ typedef struct PgStat_ShmemControl
 	 */
 	pg_atomic_uint64 gc_request_count;
 
+	/*
+	 * Counters for the number of entries associated to a single stats kind
+	 * that uses variable-numbered objects stored in the shared hash table.
+	 * These counters can be enabled on a per-kind basis, when track_counts is
+	 * set.  This counter is incremented each time a new entry is created (not
+	 * when reused), and each time an entry is dropped.
+	 */
+	pg_atomic_uint64 entry_counts[PGSTAT_KIND_MAX];
+
 	/*
 	 * Stats data for fixed-numbered objects.
 	 */
@@ -910,6 +925,17 @@ pgstat_get_entry_data(PgStat_Kind kind, PgStatShared_Common *entry)
 	return ((char *) (entry)) + off;
 }
 
+/*
+ * Returns the number of entries counted for a stats kind.
+ */
+static inline uint64
+pgstat_get_entry_count(PgStat_Kind kind)
+{
+	Assert(pgstat_get_kind_info(kind)->track_counts);
+
+	return pg_atomic_read_u64(&pgStatLocal.shmem->entry_counts[kind - 1]);
+}
+
 /*
  * Returns a pointer to the shared memory area of custom stats for
  * fixed-numbered statistics.
diff --git a/src/backend/utils/activity/pgstat.c b/src/backend/utils/activity/pgstat.c
index f8e91484e36b..f6784ad6cd15 100644
--- a/src/backend/utils/activity/pgstat.c
+++ b/src/backend/utils/activity/pgstat.c
@@ -1490,6 +1490,12 @@ pgstat_register_kind(PgStat_Kind kind, const PgStat_KindInfo *kind_info)
 			ereport(ERROR,
 					(errmsg("custom cumulative statistics property is invalid"),
 					 errhint("Custom cumulative statistics require a shared memory size for fixed-numbered objects.")));
+		if (kind_info->track_counts)
+		{
+			ereport(ERROR,
+					(errmsg("custom cumulative statistics property is invalid"),
+					 errhint("Custom cumulative statistics cannot use counter tracking for fixed-numbered objects.")));
+		}
 	}
 
 	/*
diff --git a/src/backend/utils/activity/pgstat_shmem.c b/src/backend/utils/activity/pgstat_shmem.c
index 9dc3212f7dd0..141ab10031cd 100644
--- a/src/backend/utils/activity/pgstat_shmem.c
+++ b/src/backend/utils/activity/pgstat_shmem.c
@@ -210,27 +210,35 @@ StatsShmemInit(void)
 
 		pg_atomic_init_u64(&ctl->gc_request_count, 1);
 
-		/* initialize fixed-numbered stats */
+		/* Do the per-kind initialization */
 		for (PgStat_Kind kind = PGSTAT_KIND_MIN; kind <= PGSTAT_KIND_MAX; kind++)
 		{
 			const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
 			char	   *ptr;
 
-			if (!kind_info || !kind_info->fixed_amount)
+			if (!kind_info)
 				continue;
 
-			if (pgstat_is_kind_builtin(kind))
-				ptr = ((char *) ctl) + kind_info->shared_ctl_off;
-			else
+			/* initialize counter tracking */
+			if (kind_info->track_counts)
+				pg_atomic_init_u64(&ctl->entry_counts[kind - 1], 0);
+
+			/* initialize fixed-numbered stats */
+			if (kind_info->fixed_amount)
 			{
-				int			idx = kind - PGSTAT_KIND_CUSTOM_MIN;
+				if (pgstat_is_kind_builtin(kind))
+					ptr = ((char *) ctl) + kind_info->shared_ctl_off;
+				else
+				{
+					int			idx = kind - PGSTAT_KIND_CUSTOM_MIN;
 
-				Assert(kind_info->shared_size != 0);
-				ctl->custom_data[idx] = ShmemAlloc(kind_info->shared_size);
-				ptr = ctl->custom_data[idx];
+					Assert(kind_info->shared_size != 0);
+					ctl->custom_data[idx] = ShmemAlloc(kind_info->shared_size);
+					ptr = ctl->custom_data[idx];
+				}
+
+				kind_info->init_shmem_cb(ptr);
 			}
-
-			kind_info->init_shmem_cb(ptr);
 		}
 	}
 	else
@@ -303,6 +311,7 @@ pgstat_init_entry(PgStat_Kind kind,
 	/* Create new stats entry. */
 	dsa_pointer chunk;
 	PgStatShared_Common *shheader;
+	const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
 
 	/*
 	 * Initialize refcount to 1, marking it as valid / not dropped. The entry
@@ -319,7 +328,7 @@ pgstat_init_entry(PgStat_Kind kind,
 	shhashent->dropped = false;
 
 	chunk = dsa_allocate_extended(pgStatLocal.dsa,
-								  pgstat_get_kind_info(kind)->shared_size,
+								  kind_info->shared_size,
 								  DSA_ALLOC_ZERO | DSA_ALLOC_NO_OOM);
 	if (chunk == InvalidDsaPointer)
 		return NULL;
@@ -330,6 +339,10 @@ pgstat_init_entry(PgStat_Kind kind,
 	/* Link the new entry from the hash entry. */
 	shhashent->body = chunk;
 
+	/* Increment entry count, if required. */
+	if (kind_info->track_counts)
+		pg_atomic_fetch_add_u64(&pgStatLocal.shmem->entry_counts[kind - 1], 1);
+
 	LWLockInitialize(&shheader->lock, LWTRANCHE_PGSTATS_DATA);
 
 	return shheader;
@@ -910,7 +923,17 @@ pgstat_drop_entry_internal(PgStatShared_HashEntry *shent,
 	/* release refcount marking entry as not dropped */
 	if (pg_atomic_sub_fetch_u32(&shent->refcount, 1) == 0)
 	{
+		PgStat_Kind kind = shent->key.kind;
+
 		pgstat_free_entry(shent, hstat);
+
+		/*
+		 * Entry has been dropped with refcount at 0, hence decrement the
+		 * entry counter.
+		 */
+		if (pgstat_get_kind_info(kind)->track_counts)
+			pg_atomic_sub_fetch_u64(&pgStatLocal.shmem->entry_counts[kind - 1], 1);
+
 		return true;
 	}
 	else
-- 
2.51.0

