From 173976d81c0adeeef8768d708099b0b5f4584144 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Wed, 3 Jul 2024 15:04:45 +0900
Subject: [PATCH v3 3/6] Introduce pluggable APIs for Cumulative Statistics

This commit adds support in the backend for $subject, allowing
out-of-core extensions to add their own custom statistics kinds.  The
stats kinds are divided into two parts for efficiency:
- The built-in stats kinds, with designated initializers.
- The custom kinds, able to use a range of IDs (128 slots available as
of this patch), with information saved in TopMemoryContext.

Custom cumulative statistics can only be loaded with
shared_preload_libraries at startup, and must allocate a unique ID
shared across all the PostgreSQL extension ecosystem with the following
wiki page:
https://wiki.postgresql.org/wiki/CustomCumulativeStats

As of this patch, fixed-numbered stats kind (like WAL, archiver,
bgwriter) are not supported, still some infrastructure is added to make
its support easier, as the fixed_data areas for the shmem and snapshot
structures are extended to cover custom and builtin stats kinds.
---
 src/include/pgstat.h                      |  35 ++++-
 src/include/utils/pgstat_internal.h       |   8 +-
 src/backend/utils/activity/pgstat.c       | 161 +++++++++++++++++++---
 src/backend/utils/activity/pgstat_shmem.c |   8 +-
 src/backend/utils/adt/pgstatfuncs.c       |   2 +-
 5 files changed, 186 insertions(+), 28 deletions(-)

diff --git a/src/include/pgstat.h b/src/include/pgstat.h
index 2d30fadaf1..8d523607a4 100644
--- a/src/include/pgstat.h
+++ b/src/include/pgstat.h
@@ -34,6 +34,10 @@
 /* The types of statistics entries */
 #define PgStat_Kind uint32
 
+/* Range of IDs allowed, for built-in and custom kinds */
+#define PGSTAT_KIND_MIN	1		/* Minimum ID allowed */
+#define PGSTAT_KIND_MAX	256		/* Maximum ID allowed */
+
 /* use 0 for INVALID, to catch zero-initialized data */
 #define PGSTAT_KIND_INVALID 0
 
@@ -52,9 +56,34 @@
 #define PGSTAT_KIND_SLRU	10
 #define PGSTAT_KIND_WAL	11
 
-#define PGSTAT_KIND_FIRST_VALID PGSTAT_KIND_DATABASE
-#define PGSTAT_KIND_LAST PGSTAT_KIND_WAL
-#define PGSTAT_NUM_KINDS (PGSTAT_KIND_LAST + 1)
+#define PGSTAT_KIND_MIN_BUILTIN PGSTAT_KIND_DATABASE
+#define PGSTAT_KIND_MAX_BUILTIN PGSTAT_KIND_WAL
+
+/* Custom stats kinds */
+
+/* Range of IDs allowed for custom stats kinds */
+#define PGSTAT_KIND_CUSTOM_MIN	128
+#define PGSTAT_KIND_CUSTOM_MAX	PGSTAT_KIND_MAX
+#define PGSTAT_KIND_CUSTOM_SIZE	(PGSTAT_KIND_CUSTOM_MAX - PGSTAT_KIND_CUSTOM_MIN + 1)
+
+/*
+ * PgStat_Kind to use for extensions that require an ID, but are still in
+ * development and have not reserved their own unique kind ID yet. See:
+ * https://wiki.postgresql.org/wiki/CustomCumulativeStats
+ */
+#define PGSTAT_KIND_EXPERIMENTAL	128
+
+static inline bool
+pgstat_is_kind_builtin(PgStat_Kind kind)
+{
+	return kind > PGSTAT_KIND_INVALID && kind <= PGSTAT_KIND_MAX_BUILTIN;
+}
+
+static inline bool
+pgstat_is_kind_custom(PgStat_Kind kind)
+{
+	return kind >= PGSTAT_KIND_CUSTOM_MIN && kind <= PGSTAT_KIND_CUSTOM_MAX;
+}
 
 /* Values for track_functions GUC variable --- order is significant! */
 typedef enum TrackFunctionsLevel
diff --git a/src/include/utils/pgstat_internal.h b/src/include/utils/pgstat_internal.h
index b8b2152d71..01b43a80e0 100644
--- a/src/include/utils/pgstat_internal.h
+++ b/src/include/utils/pgstat_internal.h
@@ -431,7 +431,7 @@ typedef struct PgStat_ShmemControl
 	 *
 	 * Each entry has a size of PgStat_KindInfo->shared_size.
 	 */
-	void	   *fixed_data[PGSTAT_NUM_KINDS];
+	void	   *fixed_data[PGSTAT_KIND_MAX + 1];
 
 } PgStat_ShmemControl;
 
@@ -451,8 +451,8 @@ typedef struct PgStat_Snapshot
 	 * Each entry is allocated in TopMemoryContext, for a size of
 	 * shared_data_len.
 	 */
-	bool		fixed_valid[PGSTAT_NUM_KINDS];
-	void	   *fixed_data[PGSTAT_NUM_KINDS];
+	bool		fixed_valid[PGSTAT_KIND_MAX + 1];
+	void	   *fixed_data[PGSTAT_KIND_MAX + 1];
 
 	/* to free snapshot in bulk */
 	MemoryContext context;
@@ -497,6 +497,8 @@ static inline void *pgstat_get_entry_data(PgStat_Kind kind, PgStatShared_Common
  */
 
 extern const PgStat_KindInfo *pgstat_get_kind_info(PgStat_Kind kind);
+extern void pgstat_register_kind(PgStat_Kind kind,
+								 const PgStat_KindInfo *kind_info);
 
 #ifdef USE_ASSERT_CHECKING
 extern void pgstat_assert_is_up(void);
diff --git a/src/backend/utils/activity/pgstat.c b/src/backend/utils/activity/pgstat.c
index fb27272e86..994aa2ba51 100644
--- a/src/backend/utils/activity/pgstat.c
+++ b/src/backend/utils/activity/pgstat.c
@@ -49,8 +49,16 @@
  * pgStatPending list. Pending statistics updates are flushed out by
  * pgstat_report_stat().
  *
+ * It is possible for external modules to define custom statistics kinds,
+ * that can use the same properties as any built-in stats kinds.  Each custom
+ * stats kind needs to assign a unique ID to ensure that it does not overlap
+ * with other extensions.  In order to reserve a unique stats kind ID, refer
+ * to https://wiki.postgresql.org/wiki/CustomCumulativeStats.
+ *
  * The behavior of different kinds of statistics is determined by the kind's
- * entry in pgstat_kind_infos, see PgStat_KindInfo for details.
+ * entry in pgstat_kind_builtin_infos for all the built-in statistics kinds
+ * defined, and pgstat_kind_custom_infos for custom kinds registered at
+ * startup by pgstat_register_kind().  See PgStat_KindInfo for details.
  *
  * The consistency of read accesses to statistics can be configured using the
  * stats_fetch_consistency GUC (see config.sgml and monitoring.sgml for the
@@ -253,7 +261,7 @@ static bool pgstat_is_shutdown = false;
 
 
 /*
- * The different kinds of statistics.
+ * The different kinds of built-in statistics.
  *
  * If reasonably possible, handling specific to one kind of stats should go
  * through this abstraction, rather than making more of pgstat.c aware.
@@ -265,7 +273,7 @@ static bool pgstat_is_shutdown = false;
  * seem to be a great way of doing that, given the split across multiple
  * files.
  */
-static const PgStat_KindInfo pgstat_kind_infos[PGSTAT_NUM_KINDS] = {
+static const PgStat_KindInfo pgstat_kind_builtin_infos[PGSTAT_KIND_MAX_BUILTIN + 1] = {
 
 	/* stats kinds for variable-numbered objects */
 
@@ -432,6 +440,15 @@ static const PgStat_KindInfo pgstat_kind_infos[PGSTAT_NUM_KINDS] = {
 	},
 };
 
+/*
+ * Information about custom statistics kinds.
+ *
+ * These are saved in a different array than the built-in kinds to save
+ * in clarity with the initializations.
+ *
+ * Indexed by PGSTAT_KIND_CUSTOM_MIN, of size PGSTAT_KIND_CUSTOM_SIZE.
+ */
+static const PgStat_KindInfo **pgstat_kind_custom_infos = NULL;
 
 /* ------------------------------------------------------------
  * Functions managing the state of the stats system for all backends.
@@ -997,11 +1014,11 @@ static void
 pgstat_init_snapshot(void)
 {
 	/* Initialize fixed-numbered statistics data in snapshots */
-	for (int kind = PGSTAT_KIND_FIRST_VALID; kind <= PGSTAT_KIND_LAST; kind++)
+	for (int kind = PGSTAT_KIND_MIN; kind <= PGSTAT_KIND_MAX; kind++)
 	{
 		const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
 
-		if (!kind_info->fixed_amount)
+		if (!kind_info || !kind_info->fixed_amount)
 			continue;
 
 		pgStatLocal.snapshot.fixed_data[kind] =
@@ -1102,10 +1119,12 @@ pgstat_build_snapshot(void)
 	/*
 	 * Build snapshot of all fixed-numbered stats.
 	 */
-	for (int kind = PGSTAT_KIND_FIRST_VALID; kind <= PGSTAT_KIND_LAST; kind++)
+	for (int kind = PGSTAT_KIND_MIN; kind <= PGSTAT_KIND_MAX; kind++)
 	{
 		const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
 
+		if (!kind_info)
+			continue;
 		if (!kind_info->fixed_amount)
 		{
 			Assert(kind_info->snapshot_cb == NULL);
@@ -1299,30 +1318,125 @@ pgstat_flush_pending_entries(bool nowait)
 PgStat_Kind
 pgstat_get_kind_from_str(char *kind_str)
 {
-	for (int kind = PGSTAT_KIND_FIRST_VALID; kind <= PGSTAT_KIND_LAST; kind++)
+	for (int kind = PGSTAT_KIND_MIN_BUILTIN; kind <= PGSTAT_KIND_MAX_BUILTIN; kind++)
 	{
-		if (pg_strcasecmp(kind_str, pgstat_kind_infos[kind].name) == 0)
+		if (pg_strcasecmp(kind_str, pgstat_kind_builtin_infos[kind].name) == 0)
 			return kind;
 	}
 
+	/* Check the custom set of cumulative stats */
+	if (pgstat_kind_custom_infos)
+	{
+		for (int kind = 0; kind < PGSTAT_KIND_CUSTOM_SIZE; kind++)
+		{
+			if (pgstat_kind_custom_infos[kind] &&
+				pg_strcasecmp(kind_str, pgstat_kind_custom_infos[kind]->name) == 0)
+				return kind + PGSTAT_KIND_CUSTOM_MIN;
+		}
+	}
+
 	ereport(ERROR,
 			(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 			 errmsg("invalid statistics kind: \"%s\"", kind_str)));
-	return PGSTAT_KIND_DATABASE;	/* avoid compiler warnings */
+	return PGSTAT_KIND_INVALID; /* avoid compiler warnings */
 }
 
 static inline bool
 pgstat_is_kind_valid(PgStat_Kind kind)
 {
-	return kind >= PGSTAT_KIND_FIRST_VALID && kind <= PGSTAT_KIND_LAST;
+	return pgstat_is_kind_builtin(kind) || pgstat_is_kind_custom(kind);
 }
 
 const PgStat_KindInfo *
 pgstat_get_kind_info(PgStat_Kind kind)
 {
-	Assert(pgstat_is_kind_valid(kind));
+	if (pgstat_is_kind_builtin(kind))
+		return &pgstat_kind_builtin_infos[kind];
 
-	return &pgstat_kind_infos[kind];
+	if (pgstat_is_kind_custom(kind))
+	{
+		uint32		idx = kind - PGSTAT_KIND_CUSTOM_MIN;
+
+		if (pgstat_kind_custom_infos == NULL ||
+			pgstat_kind_custom_infos[idx] == NULL)
+			return NULL;
+		return pgstat_kind_custom_infos[idx];
+	}
+
+	return NULL;
+}
+
+/*
+ * Register a new stats kind.
+ *
+ * PgStat_Kinds must be globally unique across all extensions. Refer
+ * to https://wiki.postgresql.org/wiki/CustomCumulativeStats to reserve a
+ * unique ID for your extension, to avoid conflicts with other extension
+ * developers. During development, use PGSTAT_KIND_EXPERIMENTAL to avoid
+ * needlessly reserving a new ID.
+ */
+void
+pgstat_register_kind(PgStat_Kind kind, const PgStat_KindInfo *kind_info)
+{
+	uint32		idx = kind - PGSTAT_KIND_CUSTOM_MIN;
+
+	if (kind_info->name == NULL || strlen(kind_info->name) == 0)
+		ereport(ERROR,
+				(errmsg("custom cumulative statistics name is invalid"),
+				 errhint("Provide a non-empty name for the custom cumulative statistics.")));
+
+	if (!pgstat_is_kind_custom(kind))
+		ereport(ERROR, (errmsg("custom cumulative statistics ID %u is out of range", kind),
+						errhint("Provide a custom cumulative statistics ID between %u and %u.",
+								PGSTAT_KIND_CUSTOM_MIN, PGSTAT_KIND_CUSTOM_MAX)));
+
+	if (!process_shared_preload_libraries_in_progress)
+		ereport(ERROR,
+				(errmsg("failed to register custom cumulative statistics \"%s\" with ID %u", kind_info->name, kind),
+				 errdetail("Custom cumulative statistics must be registered while initializing modules in \"shared_preload_libraries\".")));
+
+	/*
+	 * These are not supported for now, as these point out to fixed areas of
+	 * shared memory.
+	 */
+	if (kind_info->fixed_amount)
+		ereport(ERROR,
+				(errmsg("custom cumulative statistics property is invalid"),
+				 errhint("Custom cumulative statistics cannot use a fixed amount of data.")));
+
+	/*
+	 * If pgstat_kind_custom_infos is not available yet, allocate it.
+	 */
+	if (pgstat_kind_custom_infos == NULL)
+	{
+		pgstat_kind_custom_infos = (const PgStat_KindInfo **)
+			MemoryContextAllocZero(TopMemoryContext,
+								   sizeof(PgStat_KindInfo *) * PGSTAT_KIND_CUSTOM_SIZE);
+	}
+
+	if (pgstat_kind_custom_infos[idx] != NULL &&
+		pgstat_kind_custom_infos[idx]->name != NULL)
+		ereport(ERROR,
+				(errmsg("failed to register custom cumulative statistics \"%s\" with ID %u", kind_info->name, kind),
+				 errdetail("Custom cumulative statistics \"%s\" already registered with the same ID.",
+						   pgstat_kind_custom_infos[idx]->name)));
+
+	/* check for existing custom stats with the same name */
+	for (int existing_kind = 0; existing_kind < PGSTAT_KIND_CUSTOM_SIZE; existing_kind++)
+	{
+		if (pgstat_kind_custom_infos[existing_kind] == NULL)
+			continue;
+		if (!pg_strcasecmp(pgstat_kind_custom_infos[existing_kind]->name, kind_info->name))
+			ereport(ERROR,
+					(errmsg("failed to register custom cumulative statistics \"%s\" with ID %u", kind_info->name, kind),
+					 errdetail("Existing cumulative statistics with ID %u has the same name.", existing_kind + PGSTAT_KIND_CUSTOM_MIN)));
+	}
+
+	/* Register it */
+	pgstat_kind_custom_infos[idx] = kind_info;
+	ereport(LOG,
+			(errmsg("registered custom cumulative statistics \"%s\" with ID %u",
+					kind_info->name, kind)));
 }
 
 /*
@@ -1399,12 +1513,12 @@ pgstat_write_statsfile(void)
 	write_chunk_s(fpout, &format_id);
 
 	/* Write various stats structs with fixed number of objects */
-	for (int kind = PGSTAT_KIND_FIRST_VALID; kind <= PGSTAT_KIND_LAST; kind++)
+	for (int kind = PGSTAT_KIND_MIN; kind <= PGSTAT_KIND_MAX; kind++)
 	{
 		char	   *ptr;
 		const PgStat_KindInfo *info = pgstat_get_kind_info(kind);
 
-		if (!info->fixed_amount)
+		if (!info || !info->fixed_amount)
 			continue;
 
 		Assert(info->shared_size != 0 && info->shared_data_len != 0);
@@ -1434,6 +1548,17 @@ pgstat_write_statsfile(void)
 		if (ps->dropped)
 			continue;
 
+		/*
+		 * This discards data related to custom stats kinds that are unknown
+		 * to this process.
+		 */
+		if (!pgstat_is_kind_valid(ps->key.kind))
+		{
+			elog(WARNING, "found unknown stats entry %u/%u/%u",
+				 ps->key.kind, ps->key.dboid, ps->key.objoid);
+			continue;
+		}
+
 		shstats = (PgStatShared_Common *) dsa_get_address(pgStatLocal.dsa, ps->body);
 
 		kind_info = pgstat_get_kind_info(ps->key.kind);
@@ -1649,7 +1774,7 @@ pgstat_read_statsfile(void)
 					if (found)
 					{
 						dshash_release_lock(pgStatLocal.shared_hash, p);
-						elog(WARNING, "found duplicate stats entry %d/%u/%u",
+						elog(WARNING, "found duplicate stats entry %u/%u/%u",
 							 key.kind, key.dboid, key.objoid);
 						goto error;
 					}
@@ -1706,12 +1831,12 @@ pgstat_reset_after_failure(void)
 {
 	TimestampTz ts = GetCurrentTimestamp();
 
-	/* reset fixed-numbered stats */
-	for (int kind = PGSTAT_KIND_FIRST_VALID; kind <= PGSTAT_KIND_LAST; kind++)
+	/* reset fixed-numbered stats for built-in */
+	for (int kind = PGSTAT_KIND_MIN; kind <= PGSTAT_KIND_MAX; kind++)
 	{
 		const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
 
-		if (!kind_info->fixed_amount)
+		if (!kind_info || !kind_info->fixed_amount)
 			continue;
 
 		kind_info->reset_all_cb(ts);
diff --git a/src/backend/utils/activity/pgstat_shmem.c b/src/backend/utils/activity/pgstat_shmem.c
index bd62c4b72a..0b50d58cce 100644
--- a/src/backend/utils/activity/pgstat_shmem.c
+++ b/src/backend/utils/activity/pgstat_shmem.c
@@ -132,10 +132,12 @@ StatsShmemSize(void)
 	sz = add_size(sz, pgstat_dsa_init_size());
 
 	/* Add shared memory for all the fixed-numbered statistics */
-	for (int kind = PGSTAT_KIND_FIRST_VALID; kind <= PGSTAT_KIND_LAST; kind++)
+	for (int kind = PGSTAT_KIND_MIN; kind <= PGSTAT_KIND_MAX; kind++)
 	{
 		const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
 
+		if (!kind_info)
+			continue;
 		if (!kind_info->fixed_amount)
 			continue;
 
@@ -210,11 +212,11 @@ StatsShmemInit(void)
 		pg_atomic_init_u64(&ctl->gc_request_count, 1);
 
 		/* initialize fixed-numbered statistics */
-		for (int kind = PGSTAT_KIND_FIRST_VALID; kind <= PGSTAT_KIND_LAST; kind++)
+		for (int kind = PGSTAT_KIND_MIN; kind <= PGSTAT_KIND_MAX; kind++)
 		{
 			const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
 
-			if (!kind_info->fixed_amount)
+			if (!kind_info || !kind_info->fixed_amount)
 				continue;
 
 			Assert(kind_info->shared_size != 0);
diff --git a/src/backend/utils/adt/pgstatfuncs.c b/src/backend/utils/adt/pgstatfuncs.c
index 3876339ee1..3221137123 100644
--- a/src/backend/utils/adt/pgstatfuncs.c
+++ b/src/backend/utils/adt/pgstatfuncs.c
@@ -1696,7 +1696,7 @@ pg_stat_reset(PG_FUNCTION_ARGS)
  * Reset some shared cluster-wide counters
  *
  * When adding a new reset target, ideally the name should match that in
- * pgstat_kind_infos, if relevant.
+ * pgstat_kind_builtin_infos, if relevant.
  */
 Datum
 pg_stat_reset_shared(PG_FUNCTION_ARGS)
-- 
2.45.2

