From 0869c45608b7acb73afa4afc7bf2e481c239d705 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Mon, 8 Jul 2024 14:03:17 +0900
Subject: [PATCH v6 2/5] 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

This is able to support fixed-numbered (like WAL, archiver, bgwriter)
and variable-numbered stats kinds.
---
 src/include/pgstat.h                      |  35 +++-
 src/include/utils/pgstat_internal.h       |  22 ++-
 src/backend/utils/activity/pgstat.c       | 231 +++++++++++++++++++---
 src/backend/utils/activity/pgstat_shmem.c |  31 ++-
 src/backend/utils/adt/pgstatfuncs.c       |   2 +-
 5 files changed, 287 insertions(+), 34 deletions(-)

diff --git a/src/include/pgstat.h b/src/include/pgstat.h
index cc97274708..21c3705ce4 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 778f625ca1..39f63362a3 100644
--- a/src/include/utils/pgstat_internal.h
+++ b/src/include/utils/pgstat_internal.h
@@ -195,7 +195,8 @@ typedef struct PgStat_KindInfo
 
 	/*
 	 * The size of an entry in the shared stats hash table (pointed to by
-	 * PgStatShared_HashEntry->body).
+	 * PgStatShared_HashEntry->body).  For fixed-numbered statistics, this is
+	 * the size of an entry in PgStat_ShmemControl->custom_data.
 	 */
 	uint32		shared_size;
 
@@ -446,6 +447,13 @@ typedef struct PgStat_ShmemControl
 	PgStatShared_IO io;
 	PgStatShared_SLRU slru;
 	PgStatShared_Wal wal;
+
+	/*
+	 * Custom stats data with fixed-numbered objects, indexed by (PgStat_Kind
+	 * - PGSTAT_KIND_CUSTOM_MIN).
+	 */
+	void	   *custom_data[PGSTAT_KIND_CUSTOM_SIZE];
+
 } PgStat_ShmemControl;
 
 
@@ -459,7 +467,7 @@ typedef struct PgStat_Snapshot
 	/* time at which snapshot was taken */
 	TimestampTz snapshot_timestamp;
 
-	bool		fixed_valid[PGSTAT_NUM_KINDS];
+	bool		fixed_valid[PGSTAT_KIND_MAX_BUILTIN + 1];
 
 	PgStat_ArchiverStats archiver;
 
@@ -473,6 +481,14 @@ typedef struct PgStat_Snapshot
 
 	PgStat_WalStats wal;
 
+	/*
+	 * Data in snapshot for custom fixed-numbered statistics, indexed by
+	 * (PgStat_Kind - PGSTAT_KIND_CUSTOM_MIN).  Each entry is allocated in
+	 * TopMemoryContext, for a size of shared_data_len.
+	 */
+	bool		custom_valid[PGSTAT_KIND_CUSTOM_SIZE];
+	void	   *custom_data[PGSTAT_KIND_CUSTOM_SIZE];
+
 	/* to free snapshot in bulk */
 	MemoryContext context;
 	struct pgstat_snapshot_hash *stats;
@@ -516,6 +532,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 805712d414..aceb4bed00 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
@@ -174,6 +182,8 @@ typedef struct PgStat_SnapshotEntry
 static void pgstat_write_statsfile(void);
 static void pgstat_read_statsfile(void);
 
+static void pgstat_init_snapshot_fixed(void);
+
 static void pgstat_reset_after_failure(void);
 
 static bool pgstat_flush_pending_entries(bool nowait);
@@ -251,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.
@@ -263,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 */
 
@@ -436,6 +446,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.
@@ -586,6 +605,8 @@ pgstat_initialize(void)
 
 	pgstat_init_wal();
 
+	pgstat_init_snapshot_fixed();
+
 	/* Set up a process-exit hook to clean up */
 	before_shmem_exit(pgstat_shutdown_hook, 0);
 
@@ -829,6 +850,8 @@ pgstat_clear_snapshot(void)
 
 	memset(&pgStatLocal.snapshot.fixed_valid, 0,
 		   sizeof(pgStatLocal.snapshot.fixed_valid));
+	memset(&pgStatLocal.snapshot.custom_valid, 0,
+		   sizeof(pgStatLocal.snapshot.custom_valid));
 	pgStatLocal.snapshot.stats = NULL;
 	pgStatLocal.snapshot.mode = PGSTAT_FETCH_CONSISTENCY_NONE;
 
@@ -992,7 +1015,29 @@ pgstat_snapshot_fixed(PgStat_Kind kind)
 	else
 		pgstat_build_snapshot_fixed(kind);
 
-	Assert(pgStatLocal.snapshot.fixed_valid[kind]);
+	if (pgstat_is_kind_builtin(kind))
+		Assert(pgStatLocal.snapshot.fixed_valid[kind]);
+	else if (pgstat_is_kind_custom(kind))
+		Assert(pgStatLocal.snapshot.custom_valid[kind - PGSTAT_KIND_CUSTOM_MIN]);
+}
+
+static void
+pgstat_init_snapshot_fixed(void)
+{
+	/*
+	 * Initialize fixed-numbered statistics data in snapshots, only for custom
+	 * stats kinds.
+	 */
+	for (int kind = PGSTAT_KIND_CUSTOM_MIN; kind <= PGSTAT_KIND_CUSTOM_MAX; kind++)
+	{
+		const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
+
+		if (!kind_info || !kind_info->fixed_amount)
+			continue;
+
+		pgStatLocal.snapshot.custom_data[kind - PGSTAT_KIND_CUSTOM_MIN] =
+			MemoryContextAlloc(TopMemoryContext, kind_info->shared_data_len);
+	}
 }
 
 static void
@@ -1088,10 +1133,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);
@@ -1108,6 +1155,20 @@ static void
 pgstat_build_snapshot_fixed(PgStat_Kind kind)
 {
 	const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
+	int			idx;
+	bool	   *valid;
+
+	/* Position in fixed_valid or custom_valid */
+	if (pgstat_is_kind_builtin(kind))
+	{
+		idx = kind;
+		valid = pgStatLocal.snapshot.fixed_valid;
+	}
+	else
+	{
+		idx = kind - PGSTAT_KIND_CUSTOM_MIN;
+		valid = pgStatLocal.snapshot.custom_valid;
+	}
 
 	Assert(kind_info->fixed_amount);
 	Assert(kind_info->snapshot_cb != NULL);
@@ -1115,21 +1176,21 @@ pgstat_build_snapshot_fixed(PgStat_Kind kind)
 	if (pgstat_fetch_consistency == PGSTAT_FETCH_CONSISTENCY_NONE)
 	{
 		/* rebuild every time */
-		pgStatLocal.snapshot.fixed_valid[kind] = false;
+		valid[idx] = false;
 	}
-	else if (pgStatLocal.snapshot.fixed_valid[kind])
+	else if (valid[idx])
 	{
 		/* in snapshot mode we shouldn't get called again */
 		Assert(pgstat_fetch_consistency == PGSTAT_FETCH_CONSISTENCY_CACHE);
 		return;
 	}
 
-	Assert(!pgStatLocal.snapshot.fixed_valid[kind]);
+	Assert(!valid[idx]);
 
 	kind_info->snapshot_cb();
 
-	Assert(!pgStatLocal.snapshot.fixed_valid[kind]);
-	pgStatLocal.snapshot.fixed_valid[kind] = true;
+	Assert(!valid[idx]);
+	valid[idx] = true;
 }
 
 
@@ -1285,30 +1346,127 @@ 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\".")));
+
+	/*
+	 * Check some data for fixed-numbered stats.
+	 */
+	if (kind_info->fixed_amount)
+	{
+		if (kind_info->shared_size == 0)
+			ereport(ERROR,
+					(errmsg("custom cumulative statistics property is invalid"),
+					 errhint("Custom cumulative statistics require a shared memory size for fixed-numbered objects.")));
+	}
+
+	/*
+	 * 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)));
 }
 
 /*
@@ -1385,18 +1543,22 @@ pgstat_write_statsfile(void)
 	write_chunk_s(fpout, &format_id);
 
 	/* Write various stats structs for 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->snapshot_ctl_off != 0);
+		if (pgstat_is_kind_builtin(kind))
+			Assert(info->snapshot_ctl_off != 0);
 
 		pgstat_build_snapshot_fixed(kind);
-		ptr = ((char *) &pgStatLocal.snapshot) + info->snapshot_ctl_off;
+		if (pgstat_is_kind_builtin(kind))
+			ptr = ((char *) &pgStatLocal.snapshot) + info->snapshot_ctl_off;
+		else
+			ptr = pgStatLocal.snapshot.custom_data[kind - PGSTAT_KIND_CUSTOM_MIN];
 
 		fputc(PGSTAT_FILE_ENTRY_FIXED, fpout);
 		write_chunk_s(fpout, &kind);
@@ -1419,6 +1581,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);
@@ -1570,8 +1743,16 @@ pgstat_read_statsfile(void)
 						goto error;
 
 					/* Load back stats into shared memory */
-					ptr = ((char *) shmem) + info->shared_ctl_off +
-						info->shared_data_off;
+					if (pgstat_is_kind_builtin(kind))
+						ptr = ((char *) shmem) + info->shared_ctl_off +
+							info->shared_data_off;
+					else
+					{
+						int		idx = kind - PGSTAT_KIND_CUSTOM_MIN;
+
+						ptr = ((char *) shmem->custom_data[idx]) +
+							info->shared_data_off;
+					}
 
 					if (!read_chunk(fpin, ptr, info->shared_data_len))
 						goto error;
@@ -1638,7 +1819,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;
 					}
@@ -1696,11 +1877,11 @@ pgstat_reset_after_failure(void)
 	TimestampTz ts = GetCurrentTimestamp();
 
 	/* reset 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->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 1c2b69d563..ca3d3eecfb 100644
--- a/src/backend/utils/activity/pgstat_shmem.c
+++ b/src/backend/utils/activity/pgstat_shmem.c
@@ -131,6 +131,21 @@ StatsShmemSize(void)
 	sz = MAXALIGN(sizeof(PgStat_ShmemControl));
 	sz = add_size(sz, pgstat_dsa_init_size());
 
+	/* Add shared memory for all the custom fixed-numbered statistics */
+	for (int kind = PGSTAT_KIND_CUSTOM_MIN; kind <= PGSTAT_KIND_CUSTOM_MAX; kind++)
+	{
+		const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
+
+		if (!kind_info)
+			continue;
+		if (!kind_info->fixed_amount)
+			continue;
+
+		Assert(kind_info->shared_size != 0);
+
+		sz += MAXALIGN(kind_info->shared_size);
+	}
+
 	return sz;
 }
 
@@ -197,15 +212,25 @@ StatsShmemInit(void)
 		pg_atomic_init_u64(&ctl->gc_request_count, 1);
 
 		/* initialize 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);
 			char	   *ptr;
 
-			if (!kind_info->fixed_amount)
+			if (!kind_info || !kind_info->fixed_amount)
 				continue;
 
-			ptr = ((char *) ctl) + kind_info->shared_ctl_off;
+			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];
+			}
+
 			kind_info->init_shmem_cb(ptr);
 		}
 	}
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

