Add support for tracking of entry count in pgstats

Stats kinds can set a new option called "track_entry_count" (disabled by
default, available for variable-numbered stats) that will make pgstats
track the number of entries that exist in its shared hashtable.

As there is only one code path where a new entry is added, and one code
path where entries are freed, the count tracking is straight-forward in
its implementation.  Reads of these counters are optimistic, and may
change across two calls.  The counter is incremented when an entry is
created (not when reused), and is decremented when an entry is freed
from the hashtable (marked for drop with its refcount reaching 0), which
is something that pgstats decides internally.

A first use case of this facility would be pg_stat_statements, 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.  The module currently relies
on hash_get_num_entries(), which offers a cheap way to count how many
entries are in its hash table, but we cannot do that in pgstats for
variable-sized stats kinds as a single hashtable is used for all the
stats kinds.  Independently of PGSS, this is useful for other custom
stats kinds that want to cap, control, or track the number of entries
they have, without depending on a potentially expensive sequential scan
to know the number of entries while holding an extra exclusive lock.

Author: Michael Paquier <michael@paquier.xyz>
Reviewed-by: Sami Imseih <samimseih@gmail.com>
Reviewed-by: Keisuke Kuroda <keisuke.kuroda.3862@gmail.com>
Reviewed-by: Chao Li <li.evan.chao@gmail.com>
Discussion: https://postgr.es/m/aMPKWR81KT5UXvEr@paquier.xyz

src/backend/utils/activity/pgstat.c

@@ -1487,6 +1487,10 @@ 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_entry_count)
+			ereport(ERROR,
+					(errmsg("custom cumulative statistics property is invalid"),
+					 errhint("Custom cumulative statistics cannot use entry count tracking for fixed-numbered objects.")));
 	}
 
 	/*

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 entry count tracking */
+			if (kind_info->track_entry_count)
+				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_entry_count)
+		pg_atomic_fetch_add_u64(&pgStatLocal.shmem->entry_counts[kind - 1], 1);
+
 	LWLockInitialize(&shheader->lock, LWTRANCHE_PGSTATS_DATA);
 
 	return shheader;
@@ -859,6 +872,7 @@ static void
 pgstat_free_entry(PgStatShared_HashEntry *shent, dshash_seq_status *hstat)
 {
 	dsa_pointer pdsa;
+	PgStat_Kind kind = shent->key.kind;
 
 	/*
 	 * Fetch dsa pointer before deleting entry - that way we can free the
@@ -872,6 +886,10 @@ pgstat_free_entry(PgStatShared_HashEntry *shent, dshash_seq_status *hstat)
 		dshash_delete_current(hstat);
 
 	dsa_free(pgStatLocal.dsa, pdsa);
+
+	/* Decrement entry count, if required. */
+	if (pgstat_get_kind_info(kind)->track_entry_count)
+		pg_atomic_sub_fetch_u64(&pgStatLocal.shmem->entry_counts[kind - 1], 1);
 }
 
 /*

src/include/utils/pgstat_internal.h

@@ -231,6 +231,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_entry_count:1;
+
 	/*
 	 * The size of an entry in the shared stats hash table (pointed to by
 	 * PgStatShared_HashEntry->body).  For fixed-numbered statistics, this is
@@ -500,6 +506,16 @@ 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_entry_count is set.  This counter is incremented each time a new
+	 * entry is created (not reused) in the shared hash table, and is
+	 * decremented each time an entry is freed from the shared hash table.
+	 */
+	pg_atomic_uint64 entry_counts[PGSTAT_KIND_MAX];
+
 	/*
 	 * Stats data for fixed-numbered objects.
 	 */
@@ -925,6 +941,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_entry_count);
+
+	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.