v1-0002-pgstat-Add-tests-for-custom-serialization-files-a.patch

From b7b1974872ba960c1b69ecb037d098c8c2547ee0 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-46-230.ec2.internal>
Date: Wed, 19 Nov 2025 03:58:45 +0000
Subject: [PATCH v1 2/2] pgstat: Add tests for custom serialization files and
 callbacks

This commit adds recovery tests for custom stats data after a clean restart.

injection_points is enhanced to store an injection point description,
which is recorded as extra stats data.
---
 .../injection_points--1.0.sql                 |  14 +-
 .../injection_points/injection_points.c       |   8 +-
 .../injection_points/injection_stats.c        | 185 +++++++++++++++++-
 .../injection_points/injection_stats.h        |   3 +-
 src/test/recovery/t/029_stats_restart.pl      |  47 ++++-
 5 files changed, 251 insertions(+), 6 deletions(-)

diff --git a/src/test/modules/injection_points/injection_points--1.0.sql b/src/test/modules/injection_points/injection_points--1.0.sql
index a51ff538684..4861a477e98 100644
--- a/src/test/modules/injection_points/injection_points--1.0.sql
+++ b/src/test/modules/injection_points/injection_points--1.0.sql
@@ -9,10 +9,20 @@
 -- Attaches the action to the given injection point.
 --
 CREATE FUNCTION injection_points_attach(IN point_name TEXT,
-    IN action text)
+    IN action text, IN description text DEFAULT NULL)
 RETURNS void
 AS 'MODULE_PATHNAME', 'injection_points_attach'
-LANGUAGE C STRICT PARALLEL UNSAFE;
+LANGUAGE C PARALLEL UNSAFE;
+
+--
+-- injection_points_stats_description()
+--
+-- Reports statistics, if any, related to the given injection point.
+--
+CREATE FUNCTION injection_points_stats_description(IN point_name TEXT)
+RETURNS TEXT
+AS 'MODULE_PATHNAME', 'injection_points_stats_description'
+LANGUAGE C STRICT;
 
 --
 -- injection_points_attach()
diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index b7c1c58ea56..c72fb8d4cb6 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -353,9 +353,15 @@ injection_points_attach(PG_FUNCTION_ARGS)
 {
 	char	   *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
 	char	   *action = text_to_cstring(PG_GETARG_TEXT_PP(1));
+	char	   *description = NULL; /* optional description */
 	char	   *function;
 	InjectionPointCondition condition = {0};
 
+	if (PG_NARGS() > 2 && !PG_ARGISNULL(2))
+	{
+		description = text_to_cstring(PG_GETARG_TEXT_PP(2));
+	}
+
 	if (strcmp(action, "error") == 0)
 		function = "injection_error";
 	else if (strcmp(action, "notice") == 0)
@@ -386,7 +392,7 @@ injection_points_attach(PG_FUNCTION_ARGS)
 	}
 
 	/* Add entry for stats */
-	pgstat_create_inj(name);
+	pgstat_create_inj(name, description);
 
 	PG_RETURN_VOID();
 }
diff --git a/src/test/modules/injection_points/injection_stats.c b/src/test/modules/injection_points/injection_stats.c
index 158e1631af9..c70dd957f64 100644
--- a/src/test/modules/injection_points/injection_stats.c
+++ b/src/test/modules/injection_points/injection_stats.c
@@ -19,6 +19,7 @@
 #include "common/hashfn.h"
 #include "injection_stats.h"
 #include "pgstat.h"
+#include "storage/dsm_registry.h"
 #include "utils/builtins.h"
 #include "utils/pgstat_internal.h"
 
@@ -26,6 +27,7 @@
 typedef struct PgStat_StatInjEntry
 {
 	PgStat_Counter numcalls;	/* number of times point has been run */
+	dsa_pointer description;	/* injection point description */
 } PgStat_StatInjEntry;
 
 typedef struct PgStatShared_InjectionPoint
@@ -35,6 +37,14 @@ typedef struct PgStatShared_InjectionPoint
 } PgStatShared_InjectionPoint;
 
 static bool injection_stats_flush_cb(PgStat_EntryRef *entry_ref, bool nowait);
+void		injection_stats_serialize_extra(const PgStat_HashKey *key,
+											const PgStatShared_Common *header,
+											FILE *statfile,
+											FILE **extra_files);
+bool		injection_stats_deserialize_extra(PgStat_HashKey *key,
+											  const PgStatShared_Common *header,
+											  FILE *statfile,
+											  FILE **extra_files);
 
 static const PgStat_KindInfo injection_stats = {
 	.name = "injection_points",
@@ -50,6 +60,10 @@ 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_extra = injection_stats_serialize_extra,
+	.from_serialized_extra = injection_stats_deserialize_extra,
+	.num_serialized_extra_files = 2,	/* We create 2 files to test recovery,
+										 * though only one is used */
 };
 
 /*
@@ -65,6 +79,9 @@ static const PgStat_KindInfo injection_stats = {
 /* Track if stats are loaded */
 static bool inj_stats_loaded = false;
 
+/* DSA area to store an injection points description */
+dsa_area   *inj_description_dsa = NULL;
+
 /*
  * Callback for stats handling
  */
@@ -122,10 +139,13 @@ pgstat_register_inj(void)
  * Report injection point creation.
  */
 void
-pgstat_create_inj(const char *name)
+pgstat_create_inj(const char *name, const char *description)
 {
 	PgStat_EntryRef *entry_ref;
 	PgStatShared_InjectionPoint *shstatent;
+	size_t		len;
+	char	   *desc_copy;
+	bool		found;
 
 	/* leave if disabled */
 	if (!inj_stats_loaded || !inj_stats_enabled)
@@ -138,6 +158,25 @@ pgstat_create_inj(const char *name)
 
 	/* initialize shared memory data */
 	memset(&shstatent->stats, 0, sizeof(shstatent->stats));
+
+	if (!description)
+		return;
+
+	len = strlen(description) + 1;
+
+	if (!inj_description_dsa)
+		inj_description_dsa = GetNamedDSA("injection_points_description", &found);
+
+	if (inj_description_dsa)
+	{
+		shstatent->stats.description = dsa_allocate(inj_description_dsa, len);
+
+		desc_copy = dsa_get_address(inj_description_dsa,
+									shstatent->stats.description);
+		desc_copy[len - 1] = '\0';
+
+		memcpy(desc_copy, description, len);
+	}
 }
 
 /*
@@ -180,6 +219,150 @@ pgstat_report_inj(const char *name)
 	pending->numcalls++;
 }
 
+void
+injection_stats_serialize_extra(const PgStat_HashKey *key,
+								const PgStatShared_Common *header,
+								FILE *statfile,
+								FILE **extra_files)
+{
+	char	   *description;
+	size_t		qlen;
+	FILE	   *fd = extra_files[0];
+	PgStatShared_InjectionPoint *entry = (PgStatShared_InjectionPoint *) header;
+
+	/* Exit early if stats aren't available or enabled */
+	if (!inj_stats_loaded || !inj_stats_enabled || !key)
+		return;
+
+	/* Write hash key */
+	pgstat_write_chunk(fd, (void *) key, sizeof(PgStat_HashKey));
+
+	/* Handle missing description */
+	if (!DsaPointerIsValid(entry->stats.description))
+	{
+		fputc('\0', fd);
+		return;
+	}
+
+	/* Ensure DSA area is loaded */
+	if (!inj_description_dsa)
+	{
+		bool		found;
+
+		inj_description_dsa = GetNamedDSA("injection_points_description", &found);
+	}
+
+	if (!inj_description_dsa)
+	{
+		fputc('\0', fd);
+		return;
+	}
+
+	/* Get description and write it */
+	description = dsa_get_address(inj_description_dsa, entry->stats.description);
+	qlen = strlen(description) + 1; /* include null terminator */
+
+	pgstat_write_chunk(fd, description, qlen);
+}
+
+bool
+injection_stats_deserialize_extra(PgStat_HashKey *key,
+								  const PgStatShared_Common *header,
+								  FILE *statfile,
+								  FILE **extra_files)
+{
+	PgStatShared_InjectionPoint *entry;
+	dsa_pointer dp;
+	size_t		bufcap;
+	size_t		len;
+	char	   *buffer;
+	int			c;
+	FILE	   *fd = extra_files[0];
+
+	if (!inj_stats_loaded || !inj_stats_enabled)
+		return true;
+
+	/* Read the key */
+	if (!pgstat_read_chunk(fd, (void *) key, sizeof(PgStat_HashKey)))
+		return feof(fd) ? true : false;
+
+	/* Ensure DSA is ready */
+	if (!inj_description_dsa)
+	{
+		bool		found;
+
+		inj_description_dsa = GetNamedDSA("injection_points_description", &found);
+	}
+
+	entry =
+		(PgStatShared_InjectionPoint *) header;
+
+	/* Read null-terminated description */
+	bufcap = 128;
+	len = 0;
+	buffer = palloc(bufcap);
+
+	while ((c = fgetc(fd)) != EOF)
+	{
+		if (len + 1 >= bufcap)
+		{
+			bufcap *= 2;
+			buffer = repalloc(buffer, bufcap);
+		}
+
+		buffer[len++] = (char) c;
+
+		if (c == '\0')
+			break;
+	}
+
+	/* EOF reached unexpectedly */
+	if (c == EOF)
+	{
+		pfree(buffer);
+		return false;
+	}
+
+	/* Copy into DSA */
+	dp = dsa_allocate(inj_description_dsa, len);
+
+	memcpy(dsa_get_address(inj_description_dsa, dp), buffer, len);
+	entry->stats.description = dp;
+
+	pfree(buffer);
+	return true;
+}
+
+PG_FUNCTION_INFO_V1(injection_points_stats_description);
+Datum
+injection_points_stats_description(PG_FUNCTION_ARGS)
+{
+	char	   *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	PgStat_StatInjEntry *entry = pgstat_fetch_stat_injentry(name);
+
+	if (entry == NULL)
+		PG_RETURN_NULL();
+
+	if (DsaPointerIsValid(entry->description))
+	{
+		char	   *description = NULL;
+		bool		found;
+
+		if (!inj_description_dsa)
+			inj_description_dsa = GetNamedDSA("injection_points_description", &found);
+
+		if (inj_description_dsa)
+			description = dsa_get_address(inj_description_dsa, entry->description);
+
+		if (!description)
+			PG_RETURN_NULL();
+
+		PG_RETURN_TEXT_P(cstring_to_text(description));
+	}
+	else
+		PG_RETURN_NULL();
+}
+
 /*
  * SQL function returning the number of times an injection point
  * has been called.
diff --git a/src/test/modules/injection_points/injection_stats.h b/src/test/modules/injection_points/injection_stats.h
index ba310c52c7f..e50a256d91e 100644
--- a/src/test/modules/injection_points/injection_stats.h
+++ b/src/test/modules/injection_points/injection_stats.h
@@ -20,7 +20,8 @@ extern bool inj_stats_enabled;
 
 /* injection_stats.c */
 extern void pgstat_register_inj(void);
-extern void pgstat_create_inj(const char *name);
+extern void pgstat_create_inj(const char *name,
+							  const char *description);
 extern void pgstat_drop_inj(const char *name);
 extern void pgstat_report_inj(const char *name);
 
diff --git a/src/test/recovery/t/029_stats_restart.pl b/src/test/recovery/t/029_stats_restart.pl
index 021e2bf361f..1a5c9b8abad 100644
--- a/src/test/recovery/t/029_stats_restart.pl
+++ b/src/test/recovery/t/029_stats_restart.pl
@@ -10,11 +10,31 @@ use PostgreSQL::Test::Utils;
 use Test::More;
 use File::Copy;
 
+##################################################
+# This test relies on an injection point to test extra serialization data
+# is persisted after a clean shutdown.
+##################################################
+if ($ENV{enable_injection_points} ne 'yes')
+{
+	plan skip_all => 'Injection points not supported by this build';
+}
+
 my $node = PostgreSQL::Test::Cluster->new('primary');
 $node->init(allows_streaming => 1);
-$node->append_conf('postgresql.conf', "track_functions = 'all'");
+$node->append_conf('postgresql.conf',
+		"track_functions = 'all'
+		shared_preload_libraries='injection_points'
+		injection_points.stats = 'on'");
 $node->start;
 
+# Check if the extension injection_points is available, as it may be
+# possible that this script is run with installcheck, where the module
+# would not be installed by default.
+if (!$node->check_extension('injection_points'))
+{
+	plan skip_all => 'Extension injection_points not installed';
+}
+
 my $connect_db = 'postgres';
 my $db_under_test = 'test';
 
@@ -53,6 +73,14 @@ my $tableoid = $node->safe_psql($db_under_test,
 # generate stats and flush them
 trigger_funcrel_stat();
 
+# Generate custom stats data
+$node->safe_psql($connect_db, "CREATE EXTENSION injection_points");
+$node->safe_psql('postgres',
+	"SELECT injection_points_attach('stats-notice1', 'notice',
+	'this description is for stats-notice1 injection point');");
+$node->safe_psql('postgres', "SELECT injection_points_run('stats-notice1');");
+$node->safe_psql('postgres', "SELECT injection_points_run('stats-notice1');");
+
 # verify stats objects exist
 $sect = "initial";
 is(have_stats('database', $dboid, 0), 't', "$sect: db stats do exist");
@@ -70,7 +98,13 @@ ok(!-f "$statsfile", "backup statsfile cannot already exist");
 
 my $datadir = $node->data_dir();
 my $og_stats = "$datadir/pg_stat/pgstat.stat";
+my $og_custom_stats0 = "$datadir/pg_stat/pgstat.25.0.stat";
+my $og_custom_stats1 = "$datadir/pg_stat/pgstat.25.1.stat";
+my @og_tmp_stats = glob("$datadir/pg_stat/*.tmp");
+ok(!@og_tmp_stats, "origin temp stats file must exist");
 ok(-f "$og_stats", "origin stats file must exist");
+ok(-f "$og_custom_stats0", "origin custom stats file 0 must exist");
+ok(-f "$og_custom_stats1", "origin custom stats file 1 must exist");
 copy($og_stats, $statsfile) or die "Copy failed: $!";
 
 
@@ -85,9 +119,20 @@ is(have_stats('function', $dboid, $funcoid),
 is(have_stats('relation', $dboid, $tableoid),
 	't', "$sect: relation stats do exist");
 
+# Check the custom stats data survived the restart
+my $numcalls = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_numcalls('stats-notice1');");
+my $description = $node->safe_psql('postgres',
+								"SELECT injection_points_stats_description('stats-notice1');");
+is($numcalls, '2', 'number of stats calls for stats-notice1');
+is($description, 'this description is for stats-notice1 injection point',
+				 'description of stats-notice1 perserved after clean restart');
+
 $node->stop('immediate');
 
 ok(!-f "$og_stats", "no stats file should exist after immediate shutdown");
+ok(!-f "$og_custom_stats0", "no custom stats file 0 should exist after immediate shutdown");
+ok(!-f "$og_custom_stats1", "no custom stats file 1 should exist after immediate shutdown");
 
 # copy the old stats back to test we discard stats after crash restart
 copy($statsfile, $og_stats) or die "Copy failed: $!";
-- 
2.43.0