test_custom_stats: Add tests with read/write of auxiliary data

This commit builds upon 4ba012a8ed, giving an example of what can be
achieved with the new callbacks.  This provides coverage for the new
pgstats APIs, while serving as a reference template.

Note that built-in stats kinds could use them, we just don't have a
use-case there yet.

Author: Sami Imseih <samimseih@gmail.com>
Co-authored-by: Michael Paquier <michael@paquier.xyz>
Reviewed-by: Chao Li <li.evan.chao@gmail.com>
Discussion: https://postgr.es/m/CAA5RZ0s9SDOu+Z6veoJCHWk+kDeTktAtC-KY9fQ9Z6BJdDUirQ@mail.gmail.com

src/test/modules/test_custom_stats/t/001_custom_stats.pl

@@ -29,13 +29,13 @@ $node->safe_psql('postgres', q(CREATE EXTENSION test_custom_fixed_stats));
 
 # Create entries for variable-sized stats.
 $node->safe_psql('postgres',
-	q(select test_custom_stats_var_create('entry1')));
+	q(select test_custom_stats_var_create('entry1', 'Test entry 1')));
 $node->safe_psql('postgres',
-	q(select test_custom_stats_var_create('entry2')));
+	q(select test_custom_stats_var_create('entry2', 'Test entry 2')));
 $node->safe_psql('postgres',
-	q(select test_custom_stats_var_create('entry3')));
+	q(select test_custom_stats_var_create('entry3', 'Test entry 3')));
 $node->safe_psql('postgres',
-	q(select test_custom_stats_var_create('entry4')));
+	q(select test_custom_stats_var_create('entry4', 'Test entry 4')));
 
 # Update counters: entry1=2, entry2=3, entry3=2, entry4=3, fixed=3
 $node->safe_psql('postgres',
@@ -65,16 +65,28 @@ $node->safe_psql('postgres', q(select test_custom_stats_fixed_update()));
 # Test data reports.
 my $result = $node->safe_psql('postgres',
 	q(select * from test_custom_stats_var_report('entry1')));
-is($result, "entry1|2", "report for variable-sized data of entry1");
+is( $result,
+	"entry1|2|Test entry 1",
+	"report for variable-sized data of entry1");
+
 $result = $node->safe_psql('postgres',
 	q(select * from test_custom_stats_var_report('entry2')));
-is($result, "entry2|3", "report for variable-sized data of entry2");
+is( $result,
+	"entry2|3|Test entry 2",
+	"report for variable-sized data of entry2");
+
 $result = $node->safe_psql('postgres',
 	q(select * from test_custom_stats_var_report('entry3')));
-is($result, "entry3|2", "report for variable-sized data of entry3");
+is( $result,
+	"entry3|2|Test entry 3",
+	"report for variable-sized data of entry3");
+
 $result = $node->safe_psql('postgres',
 	q(select * from test_custom_stats_var_report('entry4')));
-is($result, "entry4|3", "report for variable-sized data of entry4");
+is( $result,
+	"entry4|3|Test entry 4",
+	"report for variable-sized data of entry4");
+
 $result = $node->safe_psql('postgres',
 	q(select * from test_custom_stats_fixed_report()));
 is($result, "3|", "report for fixed-sized stats");
@@ -97,7 +109,16 @@ $node->start();
 
 $result = $node->safe_psql('postgres',
 	q(select * from test_custom_stats_var_report('entry1')));
-is($result, "entry1|2", "variable-sized stats persist after clean restart");
+is( $result,
+	"entry1|2|Test entry 1",
+	"variable-sized stats persist after clean restart");
+
+$result = $node->safe_psql('postgres',
+	q(select * from test_custom_stats_var_report('entry2')));
+is( $result,
+	"entry2|3|Test entry 2",
+	"variable-sized stats persist after clean restart");
+
 $result = $node->safe_psql('postgres',
 	q(select * from test_custom_stats_fixed_report()));
 is($result, "3|", "fixed-sized stats persist after clean restart");

src/test/modules/test_custom_stats/test_custom_var_stats--1.0.sql

@@ -3,7 +3,7 @@
 -- complain if script is sourced in psql, rather than via CREATE EXTENSION
 \echo Use "CREATE EXTENSION test_custom_var_stats" to load this file. \quit
 
-CREATE FUNCTION test_custom_stats_var_create(IN name TEXT)
+CREATE FUNCTION test_custom_stats_var_create(IN name TEXT, in description TEXT)
 RETURNS void
 AS 'MODULE_PATHNAME', 'test_custom_stats_var_create'
 LANGUAGE C STRICT PARALLEL UNSAFE;
@@ -18,8 +18,9 @@ RETURNS void
 AS 'MODULE_PATHNAME', 'test_custom_stats_var_drop'
 LANGUAGE C STRICT PARALLEL UNSAFE;
 
-
-CREATE FUNCTION test_custom_stats_var_report(INOUT name TEXT, OUT calls BIGINT)
+CREATE FUNCTION test_custom_stats_var_report(INOUT name TEXT,
+                                             OUT calls BIGINT,
+                                             OUT description TEXT)
 RETURNS SETOF record
 AS 'MODULE_PATHNAME', 'test_custom_stats_var_report'
 LANGUAGE C STRICT PARALLEL UNSAFE;

src/test/modules/test_custom_stats/test_custom_var_stats.c

@@ -14,6 +14,7 @@
 
 #include "common/hashfn.h"
 #include "funcapi.h"
+#include "storage/dsm_registry.h"
 #include "utils/builtins.h"
 #include "utils/pgstat_internal.h"
 
@@ -22,6 +23,8 @@ PG_MODULE_MAGIC_EXT(
 					.version = PG_VERSION
 );
 
+#define TEST_CUSTOM_VAR_MAGIC_NUMBER (0xBEEFBEEF)
+
 /*--------------------------------------------------------------------------
  * Macros and constants
  *--------------------------------------------------------------------------
@@ -32,6 +35,9 @@ PG_MODULE_MAGIC_EXT(
  */
 #define PGSTAT_KIND_TEST_CUSTOM_VAR_STATS 25
 
+/* File paths for auxiliary data serialization */
+#define TEST_CUSTOM_AUX_DATA_DESC "pg_stat/test_custom_var_stats_desc.stats"
+
 /*
  * Hash statistic name to generate entry index for pgstat lookup.
  */
@@ -53,8 +59,23 @@ typedef struct PgStatShared_CustomVarEntry
 {
 	PgStatShared_Common header; /* standard pgstat entry header */
 	PgStat_StatCustomVarEntry stats;	/* custom statistics data */
+	dsa_pointer description;	/* pointer to description string in DSA */
 } PgStatShared_CustomVarEntry;
 
+/*--------------------------------------------------------------------------
+ * Global Variables
+ *--------------------------------------------------------------------------
+ */
+
+/* File handle for auxiliary data serialization */
+static FILE *fd_description = NULL;
+
+/* Current write offset in fd_description file */
+static pgoff_t fd_description_offset = 0;
+
+/* DSA area for storing variable-length description strings */
+static dsa_area *custom_stats_description_dsa = NULL;
+
 /*--------------------------------------------------------------------------
  * Function prototypes
  *--------------------------------------------------------------------------
@@ -64,6 +85,19 @@ typedef struct PgStatShared_CustomVarEntry
 static bool test_custom_stats_var_flush_pending_cb(PgStat_EntryRef *entry_ref,
 												   bool nowait);
 
+/* Serialization callback: write auxiliary entry data */
+static void test_custom_stats_var_to_serialized_data(const PgStat_HashKey *key,
+													 const PgStatShared_Common *header,
+													 FILE *statfile);
+
+/* Deserialization callback: read auxiliary entry data */
+static bool test_custom_stats_var_from_serialized_data(const PgStat_HashKey *key,
+													   const PgStatShared_Common *header,
+													   FILE *statfile);
+
+/* Finish callback: end of statistics file operations */
+static void test_custom_stats_var_finish(PgStat_StatsFileOp status);
+
 /*--------------------------------------------------------------------------
  * Custom kind configuration
  *--------------------------------------------------------------------------
@@ -80,6 +114,9 @@ static const PgStat_KindInfo custom_stats = {
 	.shared_data_len = sizeof(((PgStatShared_CustomVarEntry *) 0)->stats),
 	.pending_size = sizeof(PgStat_StatCustomVarEntry),
 	.flush_pending_cb = test_custom_stats_var_flush_pending_cb,
+	.to_serialized_data = test_custom_stats_var_to_serialized_data,
+	.from_serialized_data = test_custom_stats_var_from_serialized_data,
+	.finish = test_custom_stats_var_finish,
 };
 
 /*--------------------------------------------------------------------------
@@ -132,6 +169,310 @@ test_custom_stats_var_flush_pending_cb(PgStat_EntryRef *entry_ref, bool nowait)
 	return true;
 }
 
+/*
+ * test_custom_stats_var_to_serialized_data() -
+ *
+ * Serialize auxiliary data (descriptions) for custom statistics entries
+ * to a secondary statistics file. This is called while writing the statistics
+ * to disk.
+ *
+ * This callback writes a mix of data within the main pgstats file and a
+ * secondary statistics file.  The following data is written to the main file for
+ * each entry:
+ * - An arbitrary magic number.
+ * - An offset.  This is used to know the location we need to look at
+ * to retrieve the information from the second file.
+ *
+ * The following data is written to the secondary statistics file:
+ * - The entry key, cross-checked with the data from the main file
+ * when reloaded.
+ * - The length of the description.
+ * - The description data itself.
+ */
+static void
+test_custom_stats_var_to_serialized_data(const PgStat_HashKey *key,
+										 const PgStatShared_Common *header,
+										 FILE *statfile)
+{
+	char	   *description;
+	size_t		len;
+	PgStatShared_CustomVarEntry *entry = (PgStatShared_CustomVarEntry *) header;
+	bool		found;
+	uint32		magic_number = TEST_CUSTOM_VAR_MAGIC_NUMBER;
+
+	/*
+	 * First mark the main file with a magic number, keeping a trace that some
+	 * auxiliary data will exist in the secondary statistics file.
+	 */
+	pgstat_write_chunk_s(statfile, &magic_number);
+
+	/* Open statistics file for writing. */
+	if (!fd_description)
+	{
+		fd_description = AllocateFile(TEST_CUSTOM_AUX_DATA_DESC, PG_BINARY_W);
+		if (fd_description == NULL)
+		{
+			ereport(LOG,
+					(errcode_for_file_access(),
+					 errmsg("could not open statistics file \"%s\" for writing: %m",
+							TEST_CUSTOM_AUX_DATA_DESC)));
+			return;
+		}
+
+		/* Initialize offset for secondary statistics file. */
+		fd_description_offset = 0;
+	}
+
+	/* Write offset to the main data file */
+	pgstat_write_chunk_s(statfile, &fd_description_offset);
+
+	/*
+	 * First write the entry key to the secondary statistics file.  This will
+	 * be cross-checked with the key read from main stats file at loading
+	 * time.
+	 */
+	pgstat_write_chunk_s(fd_description, (PgStat_HashKey *) key);
+	fd_description_offset += sizeof(PgStat_HashKey);
+
+	if (!custom_stats_description_dsa)
+		custom_stats_description_dsa = GetNamedDSA("test_custom_stat_dsa", &found);
+
+	/* Handle entries without descriptions */
+	if (!DsaPointerIsValid(entry->description) || !custom_stats_description_dsa)
+	{
+		/* length to description file */
+		len = 0;
+		pgstat_write_chunk_s(fd_description, &len);
+		fd_description_offset += sizeof(size_t);
+		return;
+	}
+
+	/*
+	 * Retrieve description from DSA, then write the length followed by the
+	 * description.
+	 */
+	description = dsa_get_address(custom_stats_description_dsa,
+								  entry->description);
+	len = strlen(description) + 1;
+	pgstat_write_chunk_s(fd_description, &len);
+	pgstat_write_chunk(fd_description, description, len);
+
+	/*
+	 * Update offset for next entry, counting for the length (size_t) of the
+	 * description and the description contents.
+	 */
+	fd_description_offset += len + sizeof(size_t);
+}
+
+/*
+ * test_custom_stats_var_from_serialized_data() -
+ *
+ * Read auxiliary data (descriptions) for custom statistics entries from
+ * the secondary statistics file.  This is called while loading the statistics
+ * at startup.
+ *
+ * See the top of test_custom_stats_var_to_serialized_data() for a
+ * detailed description of the data layout read here.
+ */
+static bool
+test_custom_stats_var_from_serialized_data(const PgStat_HashKey *key,
+										   const PgStatShared_Common *header,
+										   FILE *statfile)
+{
+	PgStatShared_CustomVarEntry *entry;
+	dsa_pointer dp;
+	size_t		len;
+	pgoff_t		offset;
+	char	   *buffer;
+	bool		found;
+	uint32		magic_number = 0;
+	PgStat_HashKey file_key;
+
+	/* Check the magic number first, in the main file. */
+	if (!pgstat_read_chunk_s(statfile, &magic_number))
+	{
+		elog(WARNING, "failed to read magic number from statistics file");
+		return false;
+	}
+
+	if (magic_number != TEST_CUSTOM_VAR_MAGIC_NUMBER)
+	{
+		elog(WARNING, "found magic number %u from statistics file, should be %u",
+			 magic_number, TEST_CUSTOM_VAR_MAGIC_NUMBER);
+		return false;
+	}
+
+	/*
+	 * Read the offset from the main stats file, to be able to read the
+	 * auxiliary data from the secondary statistics file.
+	 */
+	if (!pgstat_read_chunk_s(statfile, &offset))
+	{
+		elog(WARNING, "failed to read metadata offset from statistics file");
+		return false;
+	}
+
+	/* Open statistics file for reading if not already open */
+	if (!fd_description)
+	{
+		fd_description = AllocateFile(TEST_CUSTOM_AUX_DATA_DESC, PG_BINARY_R);
+		if (fd_description == NULL)
+		{
+			if (errno != ENOENT)
+				ereport(LOG,
+						(errcode_for_file_access(),
+						 errmsg("could not open statistics file \"%s\" for reading: %m",
+								TEST_CUSTOM_AUX_DATA_DESC)));
+			pgstat_reset_of_kind(PGSTAT_KIND_TEST_CUSTOM_VAR_STATS);
+			return false;
+		}
+	}
+
+	/* Read data from the secondary statistics file, at the specified offset */
+	if (fseeko(fd_description, offset, SEEK_SET) != 0)
+	{
+		elog(WARNING, "failed to seek to offset %ld in description file", offset);
+		return false;
+	}
+
+	/* Read the hash key from the secondary statistics file */
+	if (!pgstat_read_chunk_s(fd_description, &file_key))
+	{
+		elog(WARNING, "failed to read hash key from file");
+		return false;
+	}
+
+	/* Check key consistency */
+	if (file_key.kind != key->kind ||
+		file_key.dboid != key->dboid ||
+		file_key.objid != key->objid)
+	{
+		elog(WARNING, "found entry key %u/%u/%" PRIu64 " not matching with %u/%u/%" PRIu64,
+			 file_key.kind, file_key.dboid, file_key.objid,
+			 key->kind, key->dboid, key->objid);
+		return false;
+	}
+
+	entry = (PgStatShared_CustomVarEntry *) header;
+
+	/* Read the description length and its data */
+	if (!pgstat_read_chunk_s(fd_description, &len))
+	{
+		elog(WARNING, "failed to read metadata length from statistics file");
+		return false;
+	}
+
+	/* Handle empty descriptions */
+	if (len == 0)
+	{
+		entry->description = InvalidDsaPointer;
+		return true;
+	}
+
+	/* Initialize DSA if needed */
+	if (!custom_stats_description_dsa)
+		custom_stats_description_dsa = GetNamedDSA("test_custom_stat_dsa", &found);
+
+	if (!custom_stats_description_dsa)
+	{
+		elog(WARNING, "could not access DSA for custom statistics descriptions");
+		return false;
+	}
+
+	buffer = palloc(len);
+	if (!pgstat_read_chunk(fd_description, buffer, len))
+	{
+		pfree(buffer);
+		elog(WARNING, "failed to read description from file");
+		return false;
+	}
+
+	/* 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);
+
+	return true;
+}
+
+/*
+ * test_custom_stats_var_finish() -
+ *
+ * Cleanup function called at the end of statistics file operations.
+ * Handles closing files and cleanup based on the operation type.
+ */
+static void
+test_custom_stats_var_finish(PgStat_StatsFileOp status)
+{
+	switch (status)
+	{
+		case STATS_WRITE:
+			if (!fd_description)
+				return;
+
+			fd_description_offset = 0;
+
+			/* Check for write errors and cleanup if necessary */
+			if (ferror(fd_description))
+			{
+				ereport(LOG,
+						(errcode_for_file_access(),
+						 errmsg("could not write to file \"%s\": %m",
+								TEST_CUSTOM_AUX_DATA_DESC)));
+				FreeFile(fd_description);
+				unlink(TEST_CUSTOM_AUX_DATA_DESC);
+			}
+			else if (FreeFile(fd_description) < 0)
+			{
+				ereport(LOG,
+						(errcode_for_file_access(),
+						 errmsg("could not close file \"%s\": %m",
+								TEST_CUSTOM_AUX_DATA_DESC)));
+				unlink(TEST_CUSTOM_AUX_DATA_DESC);
+			}
+			break;
+
+		case STATS_READ:
+			if (fd_description)
+				FreeFile(fd_description);
+
+			/* Remove the file after reading */
+			elog(DEBUG2, "removing file \"%s\"", TEST_CUSTOM_AUX_DATA_DESC);
+			unlink(TEST_CUSTOM_AUX_DATA_DESC);
+			break;
+
+		case STATS_DISCARD:
+			{
+				int			ret;
+
+				/* Attempt to remove the file */
+				ret = unlink(TEST_CUSTOM_AUX_DATA_DESC);
+				if (ret != 0)
+				{
+					if (errno == ENOENT)
+						elog(LOG,
+							 "didn't need to unlink file \"%s\" - didn't exist",
+							 TEST_CUSTOM_AUX_DATA_DESC);
+					else
+						ereport(LOG,
+								(errcode_for_file_access(),
+								 errmsg("could not unlink file \"%s\": %m",
+										TEST_CUSTOM_AUX_DATA_DESC)));
+				}
+				else
+				{
+					ereport(LOG,
+							(errmsg_internal("unlinked file \"%s\"",
+											 TEST_CUSTOM_AUX_DATA_DESC)));
+				}
+			}
+			break;
+	}
+
+	fd_description = NULL;
+}
+
 /*--------------------------------------------------------------------------
  * Helper functions
  *--------------------------------------------------------------------------
@@ -162,8 +503,7 @@ test_custom_stats_var_fetch_entry(const char *stat_name)
  * test_custom_stats_var_create
  *		Create new custom statistic entry
  *
- * Initializes a zero-valued statistics entry in shared memory.
- * Validates name length against NAMEDATALEN limit.
+ * Initializes a statistics entry with the given name and description.
  */
 PG_FUNCTION_INFO_V1(test_custom_stats_var_create);
 Datum
@@ -172,6 +512,9 @@ test_custom_stats_var_create(PG_FUNCTION_ARGS)
 	PgStat_EntryRef *entry_ref;
 	PgStatShared_CustomVarEntry *shared_entry;
 	char	   *stat_name = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	char	   *description = text_to_cstring(PG_GETARG_TEXT_PP(1));
+	dsa_pointer dp = InvalidDsaPointer;
+	bool		found;
 
 	/* Validate name length first */
 	if (strlen(stat_name) >= NAMEDATALEN)
@@ -180,6 +523,20 @@ test_custom_stats_var_create(PG_FUNCTION_ARGS)
 				 errmsg("custom statistic name \"%s\" is too long", stat_name),
 				 errdetail("Name must be less than %d characters.", NAMEDATALEN)));
 
+	/* Initialize DSA and description provided */
+	if (!custom_stats_description_dsa)
+		custom_stats_description_dsa = GetNamedDSA("test_custom_stat_dsa", &found);
+
+	if (!custom_stats_description_dsa)
+		ereport(ERROR,
+				(errmsg("could not access DSA for custom statistics descriptions")));
+
+	/* Allocate space in DSA and copy description */
+	dp = dsa_allocate(custom_stats_description_dsa, strlen(description) + 1);
+	memcpy(dsa_get_address(custom_stats_description_dsa, dp),
+		   description,
+		   strlen(description) + 1);
+
 	/* Create or get existing entry */
 	entry_ref = pgstat_get_entry_ref_locked(PGSTAT_KIND_TEST_CUSTOM_VAR_STATS, InvalidOid,
 											PGSTAT_CUSTOM_VAR_STATS_IDX(stat_name), true);
@@ -192,6 +549,9 @@ test_custom_stats_var_create(PG_FUNCTION_ARGS)
 	/* Zero-initialize statistics */
 	memset(&shared_entry->stats, 0, sizeof(shared_entry->stats));
 
+	/* Store description pointer */
+	shared_entry->description = dp;
+
 	pgstat_unlock_entry(entry_ref);
 
 	PG_RETURN_VOID();
@@ -226,8 +586,7 @@ test_custom_stats_var_update(PG_FUNCTION_ARGS)
  * test_custom_stats_var_drop
  *		Remove custom statistic entry
  *
- * Drops the named statistic from shared memory and requests
- * garbage collection if needed.
+ * Drops the named statistic from shared memory.
  */
 PG_FUNCTION_INFO_V1(test_custom_stats_var_drop);
 Datum
@@ -247,7 +606,7 @@ test_custom_stats_var_drop(PG_FUNCTION_ARGS)
  * test_custom_stats_var_report
  *		Retrieve custom statistic values
  *
- * Returns single row with statistic name and call count if the
+ * Returns single row with statistic name, call count, and description if the
  * statistic exists, otherwise returns no rows.
  */
 PG_FUNCTION_INFO_V1(test_custom_stats_var_report);
@@ -281,9 +640,13 @@ test_custom_stats_var_report(PG_FUNCTION_ARGS)
 
 	if (funcctx->call_cntr < funcctx->max_calls)
 	{
-		Datum		values[2];
-		bool		nulls[2] = {false, false};
+		Datum		values[3];
+		bool		nulls[3] = {false, false, false};
 		HeapTuple	tuple;
+		PgStat_EntryRef *entry_ref;
+		PgStatShared_CustomVarEntry *shared_entry;
+		char	   *description = NULL;
+		bool		found;
 
 		stat_name = text_to_cstring(PG_GETARG_TEXT_PP(0));
 		stat_entry = test_custom_stats_var_fetch_entry(stat_name);
@@ -291,9 +654,33 @@ test_custom_stats_var_report(PG_FUNCTION_ARGS)
 		/* Return row only if entry exists */
 		if (stat_entry)
 		{
+			/* Get entry ref to access shared entry */
+			entry_ref = pgstat_get_entry_ref(PGSTAT_KIND_TEST_CUSTOM_VAR_STATS, InvalidOid,
+											 PGSTAT_CUSTOM_VAR_STATS_IDX(stat_name), false, NULL);
+
+			if (entry_ref)
+			{
+				shared_entry = (PgStatShared_CustomVarEntry *) entry_ref->shared_stats;
+
+				/* Get description from DSA if available */
+				if (DsaPointerIsValid(shared_entry->description))
+				{
+					if (!custom_stats_description_dsa)
+						custom_stats_description_dsa = GetNamedDSA("test_custom_stat_dsa", &found);
+
+					if (custom_stats_description_dsa)
+						description = dsa_get_address(custom_stats_description_dsa, shared_entry->description);
+				}
+			}
+
 			values[0] = PointerGetDatum(cstring_to_text(stat_name));
 			values[1] = Int64GetDatum(stat_entry->numcalls);
 
+			if (description)
+				values[2] = PointerGetDatum(cstring_to_text(description));
+			else
+				nulls[2] = true;
+
 			tuple = heap_form_tuple(funcctx->tuple_desc, values, nulls);
 			SRF_RETURN_NEXT(funcctx, HeapTupleGetDatum(tuple));
 		}