v2-0002-Tests-for-custom-stat-kinds.patch

From 50fb48cfce161438ad9d3f8d39ffe0d4c7d542a4 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-46-230.ec2.internal>
Date: Tue, 2 Dec 2025 17:00:04 +0000
Subject: [PATCH v2 2/2] Tests for custom stat kinds

Creates a new test module to test custom stat
kinds. This also updates documentation to use
this module as an example for extension
developers that use custom stat kinds.
---
 doc/src/sgml/xfunc.sgml                       |   6 +-
 src/test/modules/Makefile                     |   1 +
 src/test/modules/meson.build                  |   1 +
 .../modules/test_custom_statkind/.gitignore   |   4 +
 .../modules/test_custom_statkind/Makefile     |  23 +
 .../modules/test_custom_statkind/meson.build  |  34 ++
 .../test_custom_statkind/t/001_custom_stat.pl |  92 +++
 .../test_custom_statkind--1.0.sql             |  19 +
 .../test_custom_statkind.c                    | 570 ++++++++++++++++++
 .../test_custom_statkind.control              |   4 +
 10 files changed, 753 insertions(+), 1 deletion(-)
 create mode 100644 src/test/modules/test_custom_statkind/.gitignore
 create mode 100644 src/test/modules/test_custom_statkind/Makefile
 create mode 100644 src/test/modules/test_custom_statkind/meson.build
 create mode 100644 src/test/modules/test_custom_statkind/t/001_custom_stat.pl
 create mode 100644 src/test/modules/test_custom_statkind/test_custom_statkind--1.0.sql
 create mode 100644 src/test/modules/test_custom_statkind/test_custom_statkind.c
 create mode 100644 src/test/modules/test_custom_statkind/test_custom_statkind.control

diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index 537ee6fa254..d50626523f2 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3967,6 +3967,9 @@ static const PgStat_KindInfo custom_stats = {
     .shared_data_off = offsetof(PgStatShared_Custom, stats),
     .shared_data_len = sizeof(((PgStatShared_Custom *) 0)->stats),
     .pending_size = sizeof(PgStat_StatCustomEntry),
+    .to_serialized_extra_stats = custom_stats_serialize,
+    .from_serialized_extra_stats = custom_stats_deserialize,
+    .end_extra_stats = custom_stats_file_cleanup,
 }
 </programlisting>
 
@@ -4005,7 +4008,8 @@ extern PgStat_Kind pgstat_register_kind(PgStat_Kind kind,
 
     <para>
      An example describing how to register and use custom statistics can be
-     found in <filename>src/test/modules/injection_points</filename>.
+     found in <filename>src/test/modules/test_custom_statkind</filename> and
+     <filename>src/test/modules/injection_points</filename>
     </para>
    </sect2>
 
diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index d079b91b1a2..d35e83296e6 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -21,6 +21,7 @@ SUBDIRS = \
 		  test_bloomfilter \
 		  test_copy_callbacks \
 		  test_custom_rmgrs \
+		  test_custom_statkind \
 		  test_ddl_deparse \
 		  test_dsa \
 		  test_dsm_registry \
diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build
index f5114469b92..4f19298b6ec 100644
--- a/src/test/modules/meson.build
+++ b/src/test/modules/meson.build
@@ -20,6 +20,7 @@ subdir('test_bitmapset')
 subdir('test_bloomfilter')
 subdir('test_copy_callbacks')
 subdir('test_custom_rmgrs')
+subdir('test_custom_statkind')
 subdir('test_ddl_deparse')
 subdir('test_dsa')
 subdir('test_dsm_registry')
diff --git a/src/test/modules/test_custom_statkind/.gitignore b/src/test/modules/test_custom_statkind/.gitignore
new file mode 100644
index 00000000000..5dcb3ff9723
--- /dev/null
+++ b/src/test/modules/test_custom_statkind/.gitignore
@@ -0,0 +1,4 @@
+# Generated subdirectories
+/log/
+/results/
+/tmp_check/
diff --git a/src/test/modules/test_custom_statkind/Makefile b/src/test/modules/test_custom_statkind/Makefile
new file mode 100644
index 00000000000..42f48e9c8f4
--- /dev/null
+++ b/src/test/modules/test_custom_statkind/Makefile
@@ -0,0 +1,23 @@
+# src/test/modules/test_custom_statkind/Makefile
+
+MODULE_big = test_custom_statkind
+OBJS = \
+	$(WIN32RES) \
+	test_custom_statkind.o
+PGFILEDESC = "test_custom_statkind - test code for custom stat kinds"
+
+EXTENSION = test_custom_statkind
+DATA = test_custom_statkind--1.0.sql
+
+TAP_TESTS = 1
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_custom_statkind
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_custom_statkind/meson.build b/src/test/modules/test_custom_statkind/meson.build
new file mode 100644
index 00000000000..50dc4dfd387
--- /dev/null
+++ b/src/test/modules/test_custom_statkind/meson.build
@@ -0,0 +1,34 @@
+# Copyright (c) 2025, PostgreSQL Global Development Group
+
+test_custom_statkind_sources = files(
+  'test_custom_statkind.c',
+)
+
+if host_system == 'windows'
+  test_custom_statkind_sources += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_custom_statkind',
+    '--FILEDESC', 'test_custom_statkind - test code for custom stat kinds',])
+endif
+
+test_custom_statkind = shared_module('test_custom_statkind',
+  test_custom_statkind_sources,
+  kwargs: pg_test_mod_args,
+)
+test_install_libs += test_custom_statkind
+
+test_install_data += files(
+  'test_custom_statkind.control',
+  'test_custom_statkind--1.0.sql',
+)
+
+tests += {
+  'name': 'test_custom_statkind',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'tap': {
+    'tests': [
+      't/001_custom_stat.pl',
+    ],
+    'runningcheck': false,
+  },
+}
diff --git a/src/test/modules/test_custom_statkind/t/001_custom_stat.pl b/src/test/modules/test_custom_statkind/t/001_custom_stat.pl
new file mode 100644
index 00000000000..00327407292
--- /dev/null
+++ b/src/test/modules/test_custom_statkind/t/001_custom_stat.pl
@@ -0,0 +1,92 @@
+# Copyright (c) 2021-2025, PostgreSQL Global Development Group
+
+#
+# Test custom statistics functionality
+#
+# This test verifies that custom statistics entries can be created,
+# incremented, reported, and persisted across server restarts.
+#
+
+use strict;
+use warnings FATAL => 'all';
+
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+use File::Copy;
+
+# Create a new PostgreSQL test cluster
+my $node = PostgreSQL::Test::Cluster->new('main');
+
+# Initialize the cluster and configure it to load our test module
+$node->init;
+$node->append_conf('postgresql.conf',
+	"shared_preload_libraries = test_custom_statkind");
+$node->start;
+
+# Create the extension to enable custom statistics functions
+$node->safe_psql('postgres', q(CREATE EXTENSION test_custom_statkind));
+
+# Create two custom statistics entries with descriptions
+$node->safe_psql('postgres', q(select pgstat_create_custom_stat('entry1', 'this is the description for entry1')));
+$node->safe_psql('postgres', q(select pgstat_create_custom_stat('entry2', 'this is the description for entry2')));
+
+# Increment the statistics counters:
+# entry1: 2 calls, entry2: 3 calls
+$node->safe_psql('postgres', q(select pgstat_call_custom_stat('entry1')));
+$node->safe_psql('postgres', q(select pgstat_call_custom_stat('entry1')));
+$node->safe_psql('postgres', q(select pgstat_call_custom_stat('entry2')));
+$node->safe_psql('postgres', q(select pgstat_call_custom_stat('entry2')));
+$node->safe_psql('postgres', q(select pgstat_call_custom_stat('entry2')));
+
+# Verify that the custom statistics are correctly stored and reported
+my $result = $node->safe_psql('postgres', q(select * from pgstat_report_custom_stat('entry1')));
+is($result, "entry1|this is the description for entry1|2", "entry1 should have 2 calls with description");
+
+$result = $node->safe_psql('postgres', q(select * from pgstat_report_custom_stat('entry2')));
+is($result, "entry2|this is the description for entry2|3", "entry2 should have 3 calls with description");
+
+# Test persistence across server restart
+# Perform a clean shutdown to ensure statistics are written to disk
+$node->stop();
+
+# Create a backup of the statistics file for later testing
+my $statsfile = $PostgreSQL::Test::Utils::tmp_check . '/' . "discard_custom_stats";
+ok(!-f "$statsfile", "backup statsfile should not already exist");
+
+# Locate and backup the original statistics file
+my $datadir = $node->data_dir();
+my $og_stats = "$datadir/pg_stat/test_custom_statkind.stat";
+ok(-f "$og_stats", "original stats file should exist after shutdown");
+copy($og_stats, $statsfile) or die "Copy failed: $!";
+
+# Restart the server to test statistics persistence
+$node->start();
+
+# Verify that statistics persisted across the restart
+$result = $node->safe_psql('postgres', q(select * from pgstat_report_custom_stat('entry1')));
+is($result, "entry1|this is the description for entry1|2", "entry1 stats should persist after restart");
+
+$result = $node->safe_psql('postgres', q(select * from pgstat_report_custom_stat('entry2')));
+is($result, "entry2|this is the description for entry2|3", "entry2 stats should persist after restart");
+
+# Test statistics reset behavior
+# Perform an immediate shutdown (simulates crash) to prevent stats writing
+$node->stop('immediate');
+
+# Restore the backed up statistics file and restart
+copy($statsfile, $og_stats) or die "Copy failed: $!";
+$node->start;
+
+# After immediate shutdown, the stats file should be cleaned up
+ok(!-f "$og_stats", "stats file should be removed after immediate shutdown recovery");
+
+# Verify that statistics are reset (no descriptions, zero counts)
+$result = $node->safe_psql('postgres', q(select * from pgstat_report_custom_stat('entry1')));
+is($result, "entry1||0", "entry1 should be reset after stats file cleanup");
+
+$result = $node->safe_psql('postgres', q(select * from pgstat_report_custom_stat('entry2')));
+is($result, "entry2||0", "entry2 should be reset after stats file cleanup");
+
+# Test completed successfully
+done_testing();
\ No newline at end of file
diff --git a/src/test/modules/test_custom_statkind/test_custom_statkind--1.0.sql b/src/test/modules/test_custom_statkind/test_custom_statkind--1.0.sql
new file mode 100644
index 00000000000..8b085206ac8
--- /dev/null
+++ b/src/test/modules/test_custom_statkind/test_custom_statkind--1.0.sql
@@ -0,0 +1,19 @@
+/* src/test/modules/test_custom_statkind/test_custom_statkind--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_custom_statkind" to load this file. \quit
+
+CREATE FUNCTION pgstat_create_custom_stat(IN name TEXT, IN description text)
+RETURNS void
+AS 'MODULE_PATHNAME', 'pgstat_create_custom_stat'
+LANGUAGE C STRICT PARALLEL UNSAFE;
+
+CREATE FUNCTION pgstat_call_custom_stat(IN name TEXT)
+RETURNS void
+AS 'MODULE_PATHNAME', 'pgstat_call_custom_stat'
+LANGUAGE C STRICT PARALLEL UNSAFE;
+
+CREATE FUNCTION pgstat_report_custom_stat(INOUT name TEXT , OUT description TEXT, OUT calls BIGINT)
+RETURNS SETOF record
+AS 'MODULE_PATHNAME', 'pgstat_report_custom_stat'
+LANGUAGE C STRICT PARALLEL UNSAFE;
diff --git a/src/test/modules/test_custom_statkind/test_custom_statkind.c b/src/test/modules/test_custom_statkind/test_custom_statkind.c
new file mode 100644
index 00000000000..bc9077f80a6
--- /dev/null
+++ b/src/test/modules/test_custom_statkind/test_custom_statkind.c
@@ -0,0 +1,570 @@
+/*--------------------------------------------------------------------------
+ *
+ * test_custom_statkind.c
+ *		Test module for pgstats
+ *
+ * Copyright (c) 2024-2025, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *		src/test/modules/test_custom_statkind/test_custom_statkind.c
+ *
+ * -------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include "common/hashfn.h"
+#include "funcapi.h"
+#include "pgstat.h"
+#include "storage/dsm_registry.h"
+#include "utils/builtins.h"
+#include "utils/guc.h"
+#include "utils/pgstat_internal.h"
+
+PG_MODULE_MAGIC_EXT(
+					.name = "test_custom_statkind",
+					.version = PG_VERSION
+);
+
+/*--------------------------------------------------------------------------
+ * Macros and constants
+ *--------------------------------------------------------------------------
+ */
+
+/*
+ * Kind ID reserved for test_custom_statkind. This re-uses the same ID as
+ * Injection points to avoid reserving another kind id.
+ */
+#define PGSTAT_KIND_TEST_CUSTOM_STATKIND 25
+
+/*
+ * Compute statistics entry index from point name using an 8-byte hash.
+ */
+#define PGSTAT_CUSTOM_STAT_IDX(name) hash_bytes_extended((const unsigned char *) name, strlen(name), 0)
+#define PGSTAT_CUSTOM_EXTRA_DATA_DESC "pg_stat/test_custom_statkind.stat"
+
+/*--------------------------------------------------------------------------
+ * Type definitions
+ *--------------------------------------------------------------------------
+ */
+
+/* Local statistics entry for pending data */
+typedef struct PgStat_StatCustomEntry
+{
+	PgStat_Counter numcalls;	/* number of times the entry has been looked
+								 * up */
+}			PgStat_StatCustomEntry;
+
+/* Shared memory statistics entry */
+typedef struct PgStatShared_CustomEntry
+{
+	PgStatShared_Common header;
+	PgStat_StatCustomEntry stats;
+	char		name[NAMEDATALEN];
+	dsa_pointer description;
+}			PgStatShared_CustomEntry;
+
+/*--------------------------------------------------------------------------
+ * Global variables
+ *--------------------------------------------------------------------------
+ */
+
+/* Flag indicating if custom statistics kind is loaded */
+static bool pgstat_custom_kind_loaded = false;
+
+/* File handle for statistics serialization */
+static FILE *fd = NULL;
+
+/* DSA area to store custom statistics descriptions */
+dsa_area   *custom_stats_description_dsa = NULL;
+
+/*--------------------------------------------------------------------------
+ * Function prototypes
+ *--------------------------------------------------------------------------
+ */
+
+/* Statistics callback functions */
+static bool pgstat_custom_entry_flush_cb(PgStat_EntryRef *entry_ref, bool nowait);
+static void pgstat_custom_entry_extra_serialize(const PgStat_HashKey *key,
+												const PgStatShared_Common *header, FILE *statfile);
+static void pgstat_custom_entry_extra_deserialize(const PgStat_HashKey *key,
+												  const PgStatShared_Common *header, FILE *statfile);
+static void pgstat_custom_entry_end_extra(PgStat_StatsFileOp status);
+
+/*--------------------------------------------------------------------------
+ * Custom kind configuration
+ *--------------------------------------------------------------------------
+ */
+
+static const PgStat_KindInfo custom_stats = {
+	.name = "test_custom_pgstat",
+	.fixed_amount = false,		/* Bounded by the number of points */
+	.write_to_file = true,
+	.track_entry_count = true,
+	.accessed_across_databases = true,	/* System-wide statistics */
+	.shared_size = sizeof(PgStatShared_CustomEntry),
+	.shared_data_off = offsetof(PgStatShared_CustomEntry, stats),
+	.shared_data_len = sizeof(((PgStatShared_CustomEntry *) 0)->stats),
+	.pending_size = sizeof(PgStat_StatCustomEntry),
+	.flush_pending_cb = pgstat_custom_entry_flush_cb,
+	.to_serialized_extra_stats = pgstat_custom_entry_extra_serialize,
+	.from_serialized_extra_stats = pgstat_custom_entry_extra_deserialize,
+	.end_extra_stats = pgstat_custom_entry_end_extra,
+};
+
+/*--------------------------------------------------------------------------
+ * Module initialization
+ *--------------------------------------------------------------------------
+ */
+
+void
+_PG_init(void)
+{
+	if (!process_shared_preload_libraries_in_progress)
+		return;
+
+	pgstat_register_kind(PGSTAT_KIND_TEST_CUSTOM_STATKIND, &custom_stats);
+
+	pgstat_custom_kind_loaded = true;
+}
+
+/*--------------------------------------------------------------------------
+ * Statistics callback functions
+ *--------------------------------------------------------------------------
+ */
+
+/*
+ * pgstat_custom_entry_flush_cb() -
+ *
+ * Flush callback for custom statistics entries. Merges pending statistics
+ * data from local memory into shared memory.
+ *
+ * Returns true if the flush was successful, false if we couldn't acquire
+ * the necessary locks (when nowait is true).
+ */
+static bool
+pgstat_custom_entry_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
+{
+	PgStat_StatCustomEntry *localent;
+	PgStatShared_CustomEntry *shfuncent;
+
+	localent = (PgStat_StatCustomEntry *) entry_ref->pending;
+	shfuncent = (PgStatShared_CustomEntry *) entry_ref->shared_stats;
+
+	if (!pgstat_lock_entry(entry_ref, nowait))
+		return false;
+
+	shfuncent->stats.numcalls += localent->numcalls;
+
+	pgstat_unlock_entry(entry_ref);
+
+	return true;
+}
+
+/*
+ * pgstat_custom_entry_extra_serialize() -
+ *
+ * Serialize extra data (descriptions) for custom statistics entries to
+ * the statistics file. Called during statistics file writing to preserve
+ * description strings across restarts.
+ */
+static void
+pgstat_custom_entry_extra_serialize(const PgStat_HashKey *key,
+									const PgStatShared_Common *header, FILE *statfile)
+{
+	char	   *description;
+	size_t		qlen;
+	PgStatShared_CustomEntry *entry = (PgStatShared_CustomEntry *) header;
+	bool		found;
+
+	if (!pgstat_custom_kind_loaded)
+		return;
+
+	if (!custom_stats_description_dsa)
+		custom_stats_description_dsa = GetNamedDSA("pgstat_custom_stat_dsa", &found);
+
+	/* Exit early if DSA is not available */
+	if (!custom_stats_description_dsa)
+		return;
+
+	/* Open statistics file for writing if not already open */
+	if (!fd)
+	{
+		fd = AllocateFile(PGSTAT_CUSTOM_EXTRA_DATA_DESC, PG_BINARY_W);
+		if (fd == NULL)
+		{
+			ereport(LOG,
+					(errcode_for_file_access(),
+					 errmsg("could not open statistics file \"%s\" for writing: %m",
+							PGSTAT_CUSTOM_EXTRA_DATA_DESC)));
+			return;
+		}
+	}
+
+	/* Write the hash key to identify this entry */
+	pgstat_write_chunk(fd, (void *) key, sizeof(PgStat_HashKey));
+
+	/* Handle entries without descriptions */
+	if (!DsaPointerIsValid(entry->description))
+	{
+		fputc('\0', fd);		/* Write null terminator for empty description */
+		return;
+	}
+
+	if (!custom_stats_description_dsa)
+	{
+		fputc('\0', fd);		/* Write null terminator if DSA unavailable */
+		return;
+	}
+
+	/* Retrieve description from DSA and write to file */
+	description = dsa_get_address(custom_stats_description_dsa, entry->description);
+	qlen = strlen(description) + 1; /* include null terminator */
+
+	pgstat_write_chunk(fd, description, qlen);
+}
+
+/*
+ * pgstat_custom_entry_extra_deserialize() -
+ *
+ * Deserialize extra data (descriptions) for custom statistics entries from
+ * the statistics file. Called during statistics file reading to restore
+ * description strings after a restart.
+ */
+static void
+pgstat_custom_entry_extra_deserialize(const PgStat_HashKey *key,
+									  const PgStatShared_Common *header, FILE *statfile)
+{
+	PgStatShared_CustomEntry *entry;
+	dsa_pointer dp;
+	size_t		bufsize;
+	size_t		len;
+	char	   *buffer;
+	int			c;
+	bool		found;
+
+	if (!pgstat_custom_kind_loaded)
+		return;
+
+	/* Initialize DSA if needed */
+	if (!custom_stats_description_dsa)
+		custom_stats_description_dsa = GetNamedDSA("pgstat_custom_stat_dsa", &found);
+
+	if (!custom_stats_description_dsa)
+		return;
+
+	/* Open statistics file for reading if not already open */
+	if (!fd)
+	{
+		fd = AllocateFile(PGSTAT_CUSTOM_EXTRA_DATA_DESC, PG_BINARY_R);
+		if (fd == NULL)
+		{
+			if (errno != ENOENT)
+				ereport(LOG,
+						(errcode_for_file_access(),
+						 errmsg("could not open statistics file \"%s\" for reading: %m",
+								PGSTAT_CUSTOM_EXTRA_DATA_DESC)));
+			/* Reset statistics if file is missing or unreadable */
+			pgstat_reset_of_kind(PGSTAT_KIND_TEST_CUSTOM_STATKIND);
+			return;
+		}
+	}
+
+	/* Read and verify the hash key */
+	if (!pgstat_read_chunk(fd, (void *) key, sizeof(PgStat_HashKey)))
+		return;
+
+	entry = (PgStatShared_CustomEntry *) header;
+
+	/* Read null-terminated description string from file */
+	bufsize = 128;
+	len = 0;
+	buffer = palloc(bufsize);
+
+	/* Read description character by character until null terminator */
+	while ((c = fgetc(fd)) != EOF)
+	{
+		/* Expand buffer if needed */
+		if (len + 1 >= bufsize)
+		{
+			bufsize *= 2;
+			buffer = repalloc(buffer, bufsize);
+		}
+
+		buffer[len++] = (char) c;
+
+		if (c == '\0')
+			break;
+	}
+
+	/* Handle unexpected EOF */
+	if (c == EOF)
+	{
+		pfree(buffer);
+		return;
+	}
+
+	/* Allocate space in DSA and copy the description */
+	dp = dsa_allocate(custom_stats_description_dsa, len);
+
+	memcpy(dsa_get_address(custom_stats_description_dsa, dp), buffer, len);
+	entry->description = dp;
+
+	pfree(buffer);
+}
+
+/*
+ * pgstat_custom_entry_end_extra() -
+ *
+ * Cleanup function called at the end of statistics file operations.
+ * Handles closing files and cleanup based on the operation type.
+ */
+static void
+pgstat_custom_entry_end_extra(PgStat_StatsFileOp status)
+{
+	if (!pgstat_custom_kind_loaded)
+		return;
+
+	/* Handle cleanup after writing statistics */
+	if (status == STATS_WRITE)
+	{
+		if (!fd)
+			return;
+
+		/* Check for write errors and cleanup if necessary */
+		if (ferror(fd))
+		{
+			ereport(LOG,
+					(errcode_for_file_access(),
+					 errmsg("could not write statistics file \"%s\": %m",
+							PGSTAT_CUSTOM_EXTRA_DATA_DESC)));
+			FreeFile(fd);
+			unlink(PGSTAT_CUSTOM_EXTRA_DATA_DESC);
+		}
+		else if (FreeFile(fd) < 0)
+		{
+			ereport(LOG,
+					(errcode_for_file_access(),
+					 errmsg("could not close statistics file \"%s\": %m",
+							PGSTAT_CUSTOM_EXTRA_DATA_DESC)));
+			unlink(PGSTAT_CUSTOM_EXTRA_DATA_DESC);
+		}
+	}
+	/* Handle cleanup after reading statistics */
+	else if (status == STATS_READ)
+	{
+		if (!fd)
+			return;
+
+		FreeFile(fd);
+
+		/* Remove the temporary statistics file after reading */
+		elog(DEBUG2, "removing statistics file \"%s\"", PGSTAT_CUSTOM_EXTRA_DATA_DESC);
+		unlink(PGSTAT_CUSTOM_EXTRA_DATA_DESC);
+	}
+	/* Handle other cleanup operations (reset, etc.) */
+	else
+	{
+		int			ret;
+
+		/* Attempt to remove the statistics file */
+		ret = unlink(PGSTAT_CUSTOM_EXTRA_DATA_DESC);
+		if (ret != 0)
+		{
+			if (errno == ENOENT)
+				elog(LOG,
+					 "didn't need to unlink permanent stats file \"%s\" - didn't exist",
+					 PGSTAT_CUSTOM_EXTRA_DATA_DESC);
+			else
+				ereport(LOG,
+						(errcode_for_file_access(),
+						 errmsg("could not unlink permanent statistics file \"%s\": %m",
+								PGSTAT_CUSTOM_EXTRA_DATA_DESC)));
+		}
+		else
+		{
+			ereport(LOG,
+					(errmsg_internal("unlinked permanent statistics file \"%s\"",
+									 PGSTAT_CUSTOM_EXTRA_DATA_DESC)));
+		}
+	}
+
+	fd = NULL;
+}
+
+/*--------------------------------------------------------------------------
+ * SQL-callable functions
+ *--------------------------------------------------------------------------
+ */
+
+/*
+ * pgstat_create_custom_stat() -
+ *
+ * SQL-callable function to create a new custom statistics entry.
+ * Initializes shared memory structure for tracking statistics with
+ * a name and optional description.
+ */
+PG_FUNCTION_INFO_V1(pgstat_create_custom_stat);
+Datum
+pgstat_create_custom_stat(PG_FUNCTION_ARGS)
+{
+	PgStat_EntryRef *entry_ref;
+	bool		found;
+	char	   *desc_copy;
+	size_t		len;
+	PgStatShared_CustomEntry *shstatent;
+	char	   *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	char	   *description = text_to_cstring(PG_GETARG_TEXT_PP(1));
+
+	if (!pgstat_custom_kind_loaded)
+		PG_RETURN_VOID();
+
+	entry_ref = pgstat_get_entry_ref_locked(PGSTAT_KIND_TEST_CUSTOM_STATKIND, InvalidOid,
+											PGSTAT_CUSTOM_STAT_IDX(name), true);
+
+	if (!entry_ref)
+		PG_RETURN_VOID();
+
+	shstatent = (PgStatShared_CustomEntry *) entry_ref->shared_stats;
+
+	/* Initialize shared memory statistics data to zero */
+	memset(&shstatent->stats, 0, sizeof(shstatent->stats));
+
+	/* Validate and store the statistic name */
+	if (strlen(name) >= NAMEDATALEN)
+		ereport(ERROR,
+				(errcode(ERRCODE_NAME_TOO_LONG),
+				 errmsg("custom statistic name \"%s\" is too long", name),
+				 errdetail("Name must be less than %d characters.", NAMEDATALEN)));
+
+	strcpy(shstatent->name, name);
+
+	/* Store description in DSA if provided */
+	if (description)
+	{
+		len = strlen(description) + 1;
+
+		/* Initialize DSA for descriptions if not already done */
+		if (!custom_stats_description_dsa)
+			custom_stats_description_dsa = GetNamedDSA("pgstat_custom_stat_dsa", &found);
+
+		if (custom_stats_description_dsa)
+		{
+			/* Allocate space in DSA and copy description */
+			shstatent->description = dsa_allocate(custom_stats_description_dsa, len);
+			if (!DsaPointerIsValid(shstatent->description))
+			{
+				/* DSA allocation failed, continue without description */
+				shstatent->description = InvalidDsaPointer;
+			}
+			else
+			{
+				desc_copy = dsa_get_address(custom_stats_description_dsa,
+											shstatent->description);
+				memcpy(desc_copy, description, len);
+			}
+		}
+	}
+
+	pgstat_unlock_entry(entry_ref);
+
+	PG_RETURN_VOID();
+}
+
+/*
+ * pgstat_call_custom_stat() -
+ *
+ * SQL-callable function to increment the call count for a custom statistic.
+ * This is typically called from user code to track usage or events.
+ */
+PG_FUNCTION_INFO_V1(pgstat_call_custom_stat);
+Datum
+pgstat_call_custom_stat(PG_FUNCTION_ARGS)
+{
+	char	   *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	PgStat_EntryRef *entry_ref;
+	PgStat_StatCustomEntry *pending;
+
+	if (!pgstat_custom_kind_loaded)
+		PG_RETURN_VOID();
+
+	/* Get the pending statistics entry for this custom stat */
+	entry_ref = pgstat_prep_pending_entry(PGSTAT_KIND_TEST_CUSTOM_STATKIND, InvalidOid,
+										  PGSTAT_CUSTOM_STAT_IDX(name), NULL);
+
+	pending = (PgStat_StatCustomEntry *) entry_ref->pending;
+
+	/* Increment the call counter in local pending stats */
+	pending->numcalls++;
+
+	PG_RETURN_VOID();
+}
+
+/*
+ * pgstat_report_custom_stat() -
+ *
+ * SQL-callable function to retrieve statistics for a custom statistic.
+ * Returns a composite type containing the name, description, and call count.
+ */
+PG_FUNCTION_INFO_V1(pgstat_report_custom_stat);
+Datum
+pgstat_report_custom_stat(PG_FUNCTION_ARGS)
+{
+	char	   *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	PgStat_EntryRef *entry_ref;
+	PgStatShared_CustomEntry *shfuncent;
+	Datum		values[3];
+	bool		nulls[3] = {false, false, false};
+	TupleDesc	tupdesc;
+	TupleDesc	expected;
+	bool		found;
+
+	if (!pgstat_custom_kind_loaded)
+		PG_RETURN_NULL();
+
+	/* Initialize DSA for descriptions if needed */
+	if (!custom_stats_description_dsa)
+		custom_stats_description_dsa =
+			GetNamedDSA("pgstat_custom_stat_dsa", &found);
+
+	/* Return NULL if DSA is not available */
+	if (!custom_stats_description_dsa)
+		PG_RETURN_NULL();
+
+	/* Get the return tuple descriptor from pg_proc */
+	if (get_call_result_type(fcinfo, NULL, &expected) != TYPEFUNC_COMPOSITE)
+		elog(ERROR, "pgstat_report_custom_stat: return type is not composite");
+
+	tupdesc = BlessTupleDesc(expected);
+
+	/* Look up the statistics entry in shared memory */
+	entry_ref = pgstat_get_entry_ref_locked(PGSTAT_KIND_TEST_CUSTOM_STATKIND,
+											InvalidOid,
+											PGSTAT_CUSTOM_STAT_IDX(name),
+											true);
+
+	if (entry_ref)
+	{
+		char	   *desc_copy = NULL;
+
+		shfuncent = (PgStatShared_CustomEntry *) entry_ref->shared_stats;
+
+		/* Build the return tuple with name, description, and call count */
+		values[0] = CStringGetTextDatum(name);	/* name */
+		if (DsaPointerIsValid(shfuncent->description))
+		{
+			/* Get description from DSA */
+			desc_copy = dsa_get_address(custom_stats_description_dsa,
+										shfuncent->description);
+
+			values[1] = CStringGetTextDatum(desc_copy); /* description */
+		}
+		else
+			nulls[1] = true;	/* no description available */
+
+		values[2] = Int64GetDatum(shfuncent->stats.numcalls);	/* calls */
+	}
+
+	pgstat_unlock_entry(entry_ref);
+
+	PG_RETURN_DATUM(HeapTupleGetDatum(heap_form_tuple(tupdesc, values, nulls)));
+}
diff --git a/src/test/modules/test_custom_statkind/test_custom_statkind.control b/src/test/modules/test_custom_statkind/test_custom_statkind.control
new file mode 100644
index 00000000000..7ce22cc195f
--- /dev/null
+++ b/src/test/modules/test_custom_statkind/test_custom_statkind.control
@@ -0,0 +1,4 @@
+comment = 'Test code for custom stat kinds'
+default_version = '1.0'
+module_pathname = '$libdir/test_custom_statkind'
+relocatable = true
-- 
2.43.0