From c4a26749dc69ba546c5c5a942afbcf5684189b01 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-46-230.ec2.internal>
Date: Wed, 22 Oct 2025 18:58:18 +0000
Subject: [PATCH 2/2] test custom stats callbacks

---
 .../injection_points--1.0.sql                 |  14 +-
 .../injection_points/injection_points.c       |   8 +-
 .../injection_points/injection_stats.c        | 177 +++++++++++++++++-
 .../injection_points/injection_stats.h        |   3 +-
 .../modules/injection_points/t/001_stats.pl   |  86 +++++++++
 5 files changed, 283 insertions(+), 5 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 a7b61fbdfe6..0486eea4aee 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_load()
diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index 31138301117..a2356836f7f 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..d15ef65fb39 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,12 @@ 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 *fd);
+bool		injection_stats_deserialize_extra(PgStat_HashKey *key,
+											  const PgStatShared_Common *header,
+											  FILE *fd);
 
 static const PgStat_KindInfo injection_stats = {
 	.name = "injection_points",
@@ -50,6 +58,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_extra = injection_stats_serialize_extra,
+	.from_serialized_extra = injection_stats_deserialize_extra,
 };
 
 /*
@@ -65,6 +75,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 +135,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 +154,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 +215,146 @@ pgstat_report_inj(const char *name)
 	pending->numcalls++;
 }
 
+void
+injection_stats_serialize_extra(const PgStat_HashKey *key,
+								const PgStatShared_Common *header,
+								FILE *fd)
+{
+	char	   *description;
+	size_t		qlen;
+	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 *fd)
+{
+	PgStatShared_InjectionPoint *entry;
+	dsa_pointer dp;
+	size_t		bufcap;
+	size_t		len;
+	char	   *buffer;
+	int			c;
+
+	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/modules/injection_points/t/001_stats.pl b/src/test/modules/injection_points/t/001_stats.pl
index 47ab58d0e9b..ac3961320a1 100644
--- a/src/test/modules/injection_points/t/001_stats.pl
+++ b/src/test/modules/injection_points/t/001_stats.pl
@@ -94,6 +94,92 @@ $entrycount =
   $node->safe_psql('postgres', "SELECT injection_points_stats_count();");
 is($entrycount, '0', 'number of entries after drop via SQL function');
 
+# Crash to reset stats
+$node->stop('immediate');
+$node->start;
+
+# Test custom stats serialization/de-serialization callbacks
+$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_attach('stats-notice2', 'notice',
+	'this description is for stats-notice2 injection point');");
+$node->safe_psql('postgres',
+	"SELECT injection_points_attach('stats-notice3', 'notice');");
+$node->safe_psql('postgres',
+	"SELECT injection_points_attach('stats-notice4', 'notice',
+	'this description is for stats-notice4 injection point');");
+$node->safe_psql('postgres', "SELECT injection_points_run('stats-notice1');");
+$node->safe_psql('postgres', "SELECT injection_points_run('stats-notice1');");
+$node->safe_psql('postgres', "SELECT injection_points_run('stats-notice2');");
+$node->safe_psql('postgres', "SELECT injection_points_run('stats-notice2');");
+$node->safe_psql('postgres', "SELECT injection_points_run('stats-notice3');");
+$node->safe_psql('postgres', "SELECT injection_points_run('stats-notice3');");
+$node->safe_psql('postgres', "SELECT injection_points_run('stats-notice4');");
+$node->safe_psql('postgres', "SELECT injection_points_run('stats-notice4');");
+$node->safe_psql('postgres', "SELECT injection_points_run('stats-notice4');");
+$node->safe_psql('postgres', "SELECT injection_points_run('stats-notice4');");
+$numcalls = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_numcalls('stats-notice1');");
+is($numcalls, '2', 'number of stats calls for stats-notice1');
+$numcalls = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_numcalls('stats-notice2');");
+is($numcalls, '2', 'number of stats calls for stats-notice2');
+$numcalls = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_numcalls('stats-notice3');");
+is($numcalls, '2', 'number of stats calls for stats-notice3');
+$numcalls = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_numcalls('stats-notice4');");
+is($numcalls, '4', 'number of stats calls for stats-notice4');
+my $description = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_description('stats-notice1');");
+is($description, 'this description is for stats-notice1 injection point',
+	'description of stats-notice1 perserved before clean restart');
+$description = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_description('stats-notice2');");
+is($description, 'this description is for stats-notice2 injection point',
+		 'description of stats-notice2 perserved before clean restart');
+$description = $node->safe_psql('postgres',
+				"SELECT injection_points_stats_description('stats-notice3');");
+is($description, '',
+                 'description of stats-notice3 perserved before clean restart');
+$description = $node->safe_psql('postgres',
+								"SELECT injection_points_stats_description('stats-notice4');");
+is($description, 'this description is for stats-notice4 injection point',
+                 'description of stats-notice4 perserved before clean restart');
+# clean restart
+$node->restart;
+
+$numcalls = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_numcalls('stats-notice1');");
+is($numcalls, '2', 'number of stats calls for stats-notice1');
+$numcalls = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_numcalls('stats-notice2');");
+is($numcalls, '2', 'number of stats calls for stats-notice2');
+$numcalls = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_numcalls('stats-notice3');");
+is($numcalls, '2', 'number of stats calls for stats-notice3');
+$numcalls = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_numcalls('stats-notice4');");
+is($numcalls, '4', 'number of stats calls for stats-notice4');
+$description = $node->safe_psql('postgres',
+								"SELECT injection_points_stats_description('stats-notice1');");
+is($description, 'this description is for stats-notice1 injection point',
+				 'description of stats-notice1 perserved after clean restart');
+$description = $node->safe_psql('postgres',
+								"SELECT injection_points_stats_description('stats-notice2');");
+is($description, 'this description is for stats-notice2 injection point',
+                 'description of stats-notice2 perserved after clean restart');
+$description = $node->safe_psql('postgres',
+								"SELECT injection_points_stats_description('stats-notice3');");
+is($description, '',
+				 'description of stats-notice3 perserved after clean restart');
+$description = $node->safe_psql('postgres',
+								"SELECT injection_points_stats_description('stats-notice4');");
+is($description, 'this description is for stats-notice4 injection point',
+				 'description of stats-notice4 perserved after clean restart');
+
 # Stop the server, disable the module, then restart.  The server
 # should be able to come up.
 $node->stop;
-- 
2.43.0

