From 1a096ea5bff6550815d81e80134a8e6b9aa44d6e Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Thu, 13 Jun 2024 15:54:32 +0900
Subject: [PATCH 6/6] Extend custom cumulative stats to be persistent

This patch uses the approach of a unique file, so as all the stats are
stored in the existing PGSTAT_STAT_PERMANENT_FILENAME, extending its
format with a new entry type to be able to read on load and write at
shutdown the custom kinds of statistics registered:
- The name of the custom stats kinds is stored into the stats file
rather than its ID.
- The rest of the hash key, made of the database OID and the object OID,
is stored if the custom kind is not serialized.  If serialized, where
callbacks are used to do a mapping of the on-disk name <-> hash key, the
serialized name is pushed instead.

The patch includes an example of what can be achieved with the test
module injection_points, with a TAP test checking if statistics can be
kept after a clean restart.
---
 src/include/pgstat.h                          |  2 +-
 src/backend/utils/activity/pgstat.c           | 96 ++++++++++++++++++-
 src/backend/utils/adt/pgstatfuncs.c           |  2 +-
 src/test/modules/injection_points/Makefile    |  4 +
 .../injection_points/injection_points.c       | 29 +++++-
 .../injection_points/injection_stats.c        | 54 ++++++++++-
 .../injection_points/injection_stats.h        |  4 +
 src/test/modules/injection_points/meson.build |  8 ++
 .../modules/injection_points/t/001_stats.pl   | 50 ++++++++++
 9 files changed, 240 insertions(+), 9 deletions(-)
 create mode 100644 src/test/modules/injection_points/t/001_stats.pl

diff --git a/src/include/pgstat.h b/src/include/pgstat.h
index b3cdc0da6d..6c9084b977 100644
--- a/src/include/pgstat.h
+++ b/src/include/pgstat.h
@@ -487,7 +487,7 @@ extern void pgstat_clear_snapshot(void);
 extern TimestampTz pgstat_get_stat_snapshot_timestamp(bool *have_snapshot);
 
 /* helpers */
-extern PgStat_Kind pgstat_get_kind_from_str(char *kind_str);
+extern PgStat_Kind pgstat_get_kind_from_str(char *kind_str, int elevel);
 extern bool pgstat_have_entry(PgStat_Kind kind, Oid dboid, Oid objoid);
 
 
diff --git a/src/backend/utils/activity/pgstat.c b/src/backend/utils/activity/pgstat.c
index b96743ce84..5ea7dbc64b 100644
--- a/src/backend/utils/activity/pgstat.c
+++ b/src/backend/utils/activity/pgstat.c
@@ -110,6 +110,7 @@
 #include "storage/shmem.h"
 #include "storage/spin.h"
 #include "storage/s_lock.h"
+#include "utils/builtins.h"
 #include "utils/guc_hooks.h"
 #include "utils/hsearch.h"
 #include "utils/memutils.h"
@@ -142,6 +143,7 @@
  * Identifiers in stats file.
  * ---------
  */
+#define PGSTAT_FILE_CUSTOM	'C' /* custom stats kind */
 #define PGSTAT_FILE_END		'E' /* end of file */
 #define PGSTAT_FILE_NAME	'N' /* stats entry identified by name */
 #define PGSTAT_FILE_SYSTEM	'S' /* stats entry identified by PgStat_HashKey */
@@ -1372,7 +1374,7 @@ pgstat_flush_pending_entries(bool nowait)
  */
 
 PgStat_Kind
-pgstat_get_kind_from_str(char *kind_str)
+pgstat_get_kind_from_str(char *kind_str, int elevel)
 {
 	for (int kind = PGSTAT_KIND_FIRST_VALID; kind <= PGSTAT_KIND_LAST; kind++)
 	{
@@ -1392,10 +1394,10 @@ pgstat_get_kind_from_str(char *kind_str)
 		}
 	}
 
-	ereport(ERROR,
+	ereport(elevel,
 			(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
@@ -1478,7 +1480,7 @@ pgstat_add_kind(const PgStat_KindInfo *kind_info)
 
 	if (strlen(kind_info->name) >= NAMEDATALEN)
 		elog(ERROR,
-			 "cannot use custom stats kind longer than %u characters",
+			 "cannot define custom stats kind name longer than %u characters",
 			 NAMEDATALEN - 1);
 
 	/*
@@ -1489,6 +1491,15 @@ pgstat_add_kind(const PgStat_KindInfo *kind_info)
 		elog(ERROR,
 			 "cannot define custom stats kind with fixed amount of data");
 
+	/*
+	 * Custom stats entries are represented on disk with their kind name
+	 * and their entry name, so these are mandatory.
+	 */
+	if (kind_info->to_serialized_name == NULL ||
+		kind_info->from_serialized_name == NULL)
+		elog(ERROR,
+			 "cannot define custom stats kind without serialization callbacks");
+
 	/*
 	 * Check if kind ID associated to the name is already defined, and return
 	 * it if so.
@@ -1707,7 +1718,37 @@ pgstat_write_statsfile(void)
 		/* if not dropped the valid-entry refcount should exist */
 		Assert(pg_atomic_read_u32(&ps->refcount) > 0);
 
-		if (!kind_info->to_serialized_name)
+		if (ps->key.kind > PGSTAT_KIND_LAST)
+		{
+			/*
+			 * Custom stats entry, identified by kind name and entry name on
+			 * disk.
+			 */
+			NameData	name;
+			NameData	kind_name;
+
+			fputc(PGSTAT_FILE_CUSTOM, fpout);
+
+			/* Kind name */
+			namestrcpy(&kind_name, kind_info->name);
+			write_chunk_s(fpout, &kind_name);
+
+			/*
+			 * If serialized use the on-disk format, or use the hash
+			 * key components for the database OID and object OID.
+			 */
+			if (kind_info->to_serialized_name)
+			{
+				kind_info->to_serialized_name(&ps->key, shstats, &name);
+				write_chunk_s(fpout, &name);
+			}
+			else
+			{
+				write_chunk_s(fpout, &ps->key.dboid);
+				write_chunk_s(fpout, &ps->key.objoid);
+			}
+		}
+		else if (!kind_info->to_serialized_name)
 		{
 			/* normal stats entry, identified by PgStat_HashKey */
 			fputc(PGSTAT_FILE_SYSTEM, fpout);
@@ -1874,6 +1915,7 @@ pgstat_read_statsfile(void)
 
 		switch (t)
 		{
+			case PGSTAT_FILE_CUSTOM:
 			case PGSTAT_FILE_SYSTEM:
 			case PGSTAT_FILE_NAME:
 				{
@@ -1892,6 +1934,50 @@ pgstat_read_statsfile(void)
 						if (!pgstat_is_kind_valid(key.kind))
 							goto error;
 					}
+					else if (t == PGSTAT_FILE_CUSTOM)
+					{
+						NameData	kind_name;
+						NameData	name;
+						PgStat_Kind	kind;
+						const PgStat_KindInfo *kind_info = NULL;
+
+						/* First comes the name of the stats */
+						if (!read_chunk_s(fpin, &kind_name))
+							goto error;
+
+						/* Check if it is a valid stats kind */
+						kind = pgstat_get_kind_from_str(NameStr(kind_name),
+														WARNING);
+						if (!pgstat_is_kind_valid(kind))
+							goto error;
+
+						kind_info = pgstat_get_kind_info(kind);
+
+						/* Then comes the entry name, if serialized */
+						if (kind_info->from_serialized_name)
+						{
+							if (!read_chunk_s(fpin, &name))
+								goto error;
+
+							/* Compile its key */
+							if (!kind_info->from_serialized_name(&name, &key))
+							{
+								/* skip over data for entry we don't care about */
+								if (fseek(fpin, pgstat_get_entry_len(kind), SEEK_CUR) != 0)
+									goto error;
+								continue;
+							}
+						}
+						else
+						{
+							/* Extract the rest of the hash key */
+							key.kind = kind;
+							if (!read_chunk_s(fpin, &key.dboid))
+								goto error;
+							if (!read_chunk_s(fpin, &key.objoid))
+								goto error;
+						}
+					}
 					else
 					{
 						/* stats entry identified by name on disk (e.g. slots) */
diff --git a/src/backend/utils/adt/pgstatfuncs.c b/src/backend/utils/adt/pgstatfuncs.c
index 3876339ee1..59ef463687 100644
--- a/src/backend/utils/adt/pgstatfuncs.c
+++ b/src/backend/utils/adt/pgstatfuncs.c
@@ -2028,7 +2028,7 @@ pg_stat_have_stats(PG_FUNCTION_ARGS)
 	char	   *stats_type = text_to_cstring(PG_GETARG_TEXT_P(0));
 	Oid			dboid = PG_GETARG_OID(1);
 	Oid			objoid = PG_GETARG_OID(2);
-	PgStat_Kind kind = pgstat_get_kind_from_str(stats_type);
+	PgStat_Kind kind = pgstat_get_kind_from_str(stats_type, ERROR);
 
 	PG_RETURN_BOOL(pgstat_have_entry(kind, dboid, objoid));
 }
diff --git a/src/test/modules/injection_points/Makefile b/src/test/modules/injection_points/Makefile
index 676823a87f..9d84d9aaf3 100644
--- a/src/test/modules/injection_points/Makefile
+++ b/src/test/modules/injection_points/Makefile
@@ -12,9 +12,13 @@ PGFILEDESC = "injection_points - facility for injection points"
 REGRESS = injection_points
 REGRESS_OPTS = --dlpath=$(top_builddir)/src/test/regress
 
+TAP_TESTS = 1
+
 # The injection points are cluster-wide, so disable installcheck
 NO_INSTALLCHECK = 1
 
+export enable_injection_points enable_injection_points
+
 ifdef USE_PGXS
 PG_CONFIG = pg_config
 PGXS := $(shell $(PG_CONFIG) --pgxs)
diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index 6346af6cf6..f0c8c67826 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -34,9 +34,11 @@
 
 PG_MODULE_MAGIC;
 
+/* Hooks */
+static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
+
 /* Maximum number of waits usable in injection points at once */
 #define INJ_MAX_WAIT	8
-#define INJ_NAME_MAXLEN	64
 
 /*
  * Conditions related to injection points.  This tracks in shared memory the
@@ -419,3 +421,28 @@ injection_points_detach(PG_FUNCTION_ARGS)
 
 	PG_RETURN_VOID();
 }
+
+static void
+injection_points_shmem_startup(void)
+{
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	/*
+	 * Note that this does not call injection_init_shmem(), as it relies
+	 * on the DSM registry, which cannot be used in the postmaster context.
+	 */
+
+	/* Register custom statistics */
+	pgstat_register_inj();
+}
+
+void
+_PG_init(void)
+{
+	if (!process_shared_preload_libraries_in_progress)
+		return;
+
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = injection_points_shmem_startup;
+}
diff --git a/src/test/modules/injection_points/injection_stats.c b/src/test/modules/injection_points/injection_stats.c
index 6314486d79..3acbccd32a 100644
--- a/src/test/modules/injection_points/injection_stats.c
+++ b/src/test/modules/injection_points/injection_stats.c
@@ -22,9 +22,13 @@
 #include "utils/builtins.h"
 #include "utils/pgstat_internal.h"
 
+/* Maximum name length */
+#define INJ_NAME_MAXLEN 64
+
 /* Structures for statistics of injection points */
 typedef struct PgStat_StatInjEntry
 {
+	char		name[INJ_NAME_MAXLEN];	/* used for serialization */
 	PgStat_Counter numcalls;	/* number of times point has been run */
 } PgStat_StatInjEntry;
 
@@ -35,6 +39,11 @@ typedef struct PgStatShared_InjectionPoint
 } PgStatShared_InjectionPoint;
 
 static bool injection_stats_flush_cb(PgStat_EntryRef *entry_ref, bool nowait);
+static void injection_stats_to_serialized_name_cb(const PgStat_HashKey *key,
+												  const PgStatShared_Common *header,
+												  NameData *name);
+static bool injection_stats_from_serialized_name_cb(const NameData *name,
+													PgStat_HashKey *key);
 
 static const PgStat_KindInfo injection_stats = {
 	.name = "injection_points",
@@ -48,6 +57,8 @@ static const PgStat_KindInfo injection_stats = {
 	.shared_data_len = sizeof(((PgStatShared_InjectionPoint *) 0)->stats),
 	.pending_size = sizeof(PgStat_StatInjEntry),
 	.flush_pending_cb = injection_stats_flush_cb,
+	.to_serialized_name = injection_stats_to_serialized_name_cb,
+	.from_serialized_name = injection_stats_from_serialized_name_cb,
 };
 
 /*
@@ -58,7 +69,7 @@ static const PgStat_KindInfo injection_stats = {
 static PgStat_Kind inj_stats_kind = PGSTAT_KIND_INVALID;
 
 /*
- * Callback for stats handling
+ * Callbacks for stats handling
  */
 static bool
 injection_stats_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
@@ -76,6 +87,35 @@ injection_stats_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 	return true;
 }
 
+/*
+ * Transforms a hash key into a name that is then stored on disk.
+ */
+static void
+injection_stats_to_serialized_name_cb(const PgStat_HashKey *key,
+									  const PgStatShared_Common *header,
+									  NameData *name)
+{
+	PgStatShared_InjectionPoint *shstatent =
+		(PgStatShared_InjectionPoint *) header;
+
+	namestrcpy(name, shstatent->stats.name);
+}
+
+/*
+ * Computes the hash key used for this injection point name.
+ */
+static bool
+injection_stats_from_serialized_name_cb(const NameData *name,
+										PgStat_HashKey *key)
+{
+	uint32		idx = PGSTAT_INJ_IDX(NameStr(*name));
+
+	key->kind = inj_stats_kind;
+	key->dboid = InvalidOid;
+	key->objoid = idx;
+	return true;
+}
+
 /*
  * Support function for the SQL-callable pgstat* functions.  Returns
  * a pointer to the injection point statistics struct.
@@ -95,6 +135,16 @@ pgstat_fetch_stat_injentry(const char *name)
 	return entry;
 }
 
+/*
+ * Workhorse to do the registration work, called in _PG_init().
+ */
+void
+pgstat_register_inj(void)
+{
+	if (inj_stats_kind == PGSTAT_KIND_INVALID)
+		inj_stats_kind = pgstat_add_kind(&injection_stats);
+}
+
 /*
  * Report injection point creation.
  */
@@ -113,6 +163,8 @@ pgstat_create_inj(const char *name)
 
 	/* initialize shared memory data */
 	memset(&shstatent->stats, 0, sizeof(shstatent->stats));
+	strlcpy(shstatent->stats.name, name, INJ_NAME_MAXLEN);
+
 	pgstat_unlock_entry(entry_ref);
 }
 
diff --git a/src/test/modules/injection_points/injection_stats.h b/src/test/modules/injection_points/injection_stats.h
index 15621f49fd..b59ac4cf1c 100644
--- a/src/test/modules/injection_points/injection_stats.h
+++ b/src/test/modules/injection_points/injection_stats.h
@@ -15,6 +15,10 @@
 #ifndef INJECTION_STATS
 #define INJECTION_STATS
 
+/* Maximum name length */
+#define INJ_NAME_MAXLEN 64
+
+extern void pgstat_register_inj(void);
 extern void pgstat_create_inj(const char *name);
 extern void pgstat_drop_inj(const char *name);
 extern void pgstat_report_inj(const char *name);
diff --git a/src/test/modules/injection_points/meson.build b/src/test/modules/injection_points/meson.build
index 526fbc1457..612216c144 100644
--- a/src/test/modules/injection_points/meson.build
+++ b/src/test/modules/injection_points/meson.build
@@ -38,4 +38,12 @@ tests += {
     # The injection points are cluster-wide, so disable installcheck
     'runningcheck': false,
   },
+  'tap': {
+    'env': {
+      'enable_injection_points': get_option('injection_points') ? 'yes' : 'no',
+    },
+    'tests': [
+      't/001_stats.pl',
+    ],
+  },
 }
diff --git a/src/test/modules/injection_points/t/001_stats.pl b/src/test/modules/injection_points/t/001_stats.pl
new file mode 100644
index 0000000000..f8d90a3869
--- /dev/null
+++ b/src/test/modules/injection_points/t/001_stats.pl
@@ -0,0 +1,50 @@
+
+# Copyright (c) 2024, PostgreSQL Global Development Group
+
+use strict;
+use warnings FATAL => 'all';
+use locale;
+
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+# Test persistency of statistics generated for injection points.
+if ($ENV{enable_injection_points} ne 'yes')
+{
+	plan skip_all => 'Injection points not supported by this build';
+}
+
+# Node initialization
+my $node = PostgreSQL::Test::Cluster->new('master');
+$node->init;
+$node->append_conf('postgresql.conf',
+	"shared_preload_libraries = 'injection_points'");
+$node->start;
+$node->safe_psql('postgres', 'CREATE EXTENSION injection_points;');
+
+# This should count for two calls.
+$node->safe_psql('postgres',
+	"SELECT injection_points_attach('stats-notice', 'notice');");
+$node->safe_psql('postgres',
+	"SELECT injection_points_run('stats-notice');");
+$node->safe_psql('postgres',
+	"SELECT injection_points_run('stats-notice');");
+my $numcalls = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_numcalls('stats-notice');");
+is($numcalls, '2', 'number of stats calls');
+
+# Restart the node cleanly, stats should still be around.
+$node->restart;
+$numcalls = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_numcalls('stats-notice');");
+is($numcalls, '2', 'number of stats after clean restart');
+
+# On crash the stats are gone.
+$node->stop('immediate');
+$node->start;
+$numcalls = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_numcalls('stats-notice');");
+is($numcalls, '', 'number of stats after clean restart');
+
+done_testing();
-- 
2.43.0

