0002-Testing-module.patch

text/plain

Filename: 0002-Testing-module.patch
Type: text/plain
Part: 1
Message: Re: Comments on Custom RMGRs
From f92abbcc3667103628608d248870867200087e16 Mon Sep 17 00:00:00 2001
From: "Andrei V. Lepikhov" <lepihov@gmail.com>
Date: Fri, 14 Nov 2025 16:35:21 +0100
Subject: [PATCH 2/2] Testing module

---
 src/test/modules/test_dsm_registry/Makefile   |   1 +
 .../test_dsm_registry/t/001_file_storage.pl   |  31 ++++
 .../test_dsm_registry/test_dsm_registry.c     | 163 ++++++++++++++++++
 3 files changed, 195 insertions(+)
 create mode 100644 src/test/modules/test_dsm_registry/t/001_file_storage.pl

diff --git a/src/test/modules/test_dsm_registry/Makefile b/src/test/modules/test_dsm_registry/Makefile
index b13e99a354f..9aae8b98aba 100644
--- a/src/test/modules/test_dsm_registry/Makefile
+++ b/src/test/modules/test_dsm_registry/Makefile
@@ -10,6 +10,7 @@ EXTENSION = test_dsm_registry
 DATA = test_dsm_registry--1.0.sql
 
 REGRESS = test_dsm_registry
+TAP_TESTS = 1
 
 ifdef USE_PGXS
 PG_CONFIG = pg_config
diff --git a/src/test/modules/test_dsm_registry/t/001_file_storage.pl b/src/test/modules/test_dsm_registry/t/001_file_storage.pl
new file mode 100644
index 00000000000..0e82d0adcf7
--- /dev/null
+++ b/src/test/modules/test_dsm_registry/t/001_file_storage.pl
@@ -0,0 +1,31 @@
+# Copyright (c) 2023-2025, PostgreSQL Global Development Group
+use strict;
+use warnings FATAL => 'all';
+use Config;
+use PostgreSQL::Test::Utils;
+use PostgreSQL::Test::Cluster;
+use Test::More;
+
+my $node = PostgreSQL::Test::Cluster->new('node');
+
+$node->init();
+$node->append_conf('postgresql.conf',
+							"shared_preload_libraries = 'test_dsm_registry'");
+$node->start();
+
+$node->safe_psql('postgres', "CREATE EXTENSION test_dsm_registry");
+
+my $result;
+
+$node->safe_psql('postgres', "SELECT set_val_in_hash('test-1', '1414')");
+$node->safe_psql('postgres', 'CHECKPOINT');
+$node->safe_psql('postgres', "SELECT set_val_in_hash('test-2', '1415')");
+$node->stop('immediate');
+$node->start();
+
+$result = $node->safe_psql('postgres', "SELECT get_val_in_hash('test-1')");
+is($result, '1414', "Value inserted before the checkpoint was restored");
+$result = $node->safe_psql('postgres', "SELECT get_val_in_hash('test-2')");
+is($result, '', "Value inserted after the checkpoint was lost");
+
+done_testing();
diff --git a/src/test/modules/test_dsm_registry/test_dsm_registry.c b/src/test/modules/test_dsm_registry/test_dsm_registry.c
index 4cc2ccdac3f..2d7fd35a74d 100644
--- a/src/test/modules/test_dsm_registry/test_dsm_registry.c
+++ b/src/test/modules/test_dsm_registry/test_dsm_registry.c
@@ -12,13 +12,22 @@
  */
 #include "postgres.h"
 
+#include "access/xlog.h"
 #include "fmgr.h"
+#include "pgstat.h"
 #include "storage/dsm_registry.h"
+#include "storage/fd.h"
 #include "storage/lwlock.h"
 #include "utils/builtins.h"
+#include "utils/hsearch.h"
 
 PG_MODULE_MAGIC;
 
+/* Location of permanent storage file (valid on checkpoint) */
+#define TDR_DUMP_FILE	PGSTAT_STAT_PERMANENT_DIRECTORY "/pg_stat_statements.stat"
+/* Magic number identifying the stats file format */
+static const uint32 TDR_FILE_HEADER = 0x20251114;
+
 typedef struct TestDSMRegistryStruct
 {
 	int			val;
@@ -43,6 +52,11 @@ static const dshash_parameters dsh_params = {
 	dshash_strcpy
 };
 
+static Checkpoint_hook_type	prev_Checkpoint_hook = NULL;
+
+static void load_htab(void);
+static void pgss_Checkpoint(XLogRecPtr checkPointRedo, int flags);
+
 static void
 init_tdr_dsm(void *ptr)
 {
@@ -66,7 +80,14 @@ tdr_attach_shmem(void)
 		tdr_dsa = GetNamedDSA("test_dsm_registry_dsa", &found);
 
 	if (tdr_hash == NULL)
+	{
+		LWLockAcquire(&tdr_dsm->lck, LW_EXCLUSIVE);
 		tdr_hash = GetNamedDSHash("test_dsm_registry_hash", &dsh_params, &found);
+		if (!found)
+			load_htab();
+
+		LWLockRelease(&tdr_dsm->lck);
+	}
 }
 
 PG_FUNCTION_INFO_V1(set_val_in_shmem);
@@ -144,3 +165,145 @@ get_val_in_hash(PG_FUNCTION_ARGS)
 
 	PG_RETURN_TEXT_P(val);
 }
+
+/*
+ * Load any pre-existing entries from file.
+ */
+static void
+load_htab(void)
+{
+	bool	found;
+	FILE   *file = NULL;
+	uint32	header;
+	char   *val = palloc(1);
+
+	Assert(tdr_dsa != NULL && tdr_hash != NULL);
+
+	/*
+	 * Attempt to load old entries from the dump file.
+	 */
+	file = AllocateFile(TDR_DUMP_FILE, PG_BINARY_R);
+	if (file == NULL)
+	{
+		if (errno != ENOENT)
+			goto read_error;
+		/* No existing persisted file, so we're done */
+		return;
+	}
+
+	if (fread(&header, sizeof(uint32), 1, file) != 1 ||
+		header != TDR_FILE_HEADER)
+		goto read_error;
+
+	while (!feof(file))
+	{
+		TestDSMRegistryHashEntry *entry;
+		char	key[64];
+		int		keylen = offsetof(TestDSMRegistryHashEntry, val);
+		int32	vlen;
+
+		if (fread(key, keylen, 1, file) != 1 ||
+			fread(&vlen, sizeof(int32), 1, file) != 1)
+			goto read_error;
+
+		val = repalloc(val, vlen);
+		if (fread(val, vlen, 1, file) != 1)
+			goto read_error;
+
+		Assert(val[vlen - 1] == '\0');
+
+		entry = (TestDSMRegistryHashEntry *)
+								dshash_find_or_insert(tdr_hash, key, &found);
+		Assert(!found);
+
+		entry->val = dsa_allocate(tdr_dsa, strlen(val) + 1);
+		strcpy(dsa_get_address(tdr_dsa, entry->val), val);
+
+		dshash_release_lock(tdr_hash, entry);
+	}
+
+	FreeFile(file);
+	return;
+
+read_error:
+	ereport(LOG,
+			(errcode_for_file_access(),
+			 errmsg("could not read from file \"%s\": %m", TDR_DUMP_FILE)));
+	if (file)
+		FreeFile(file);
+	/* If possible, throw away the bogus file; ignore any error */
+	unlink(TDR_DUMP_FILE);
+}
+
+/*
+ * Dump hash table into file.
+ *
+ */
+static void
+pgss_Checkpoint(XLogRecPtr checkPointRedo, int flags)
+{
+	FILE					   *file;
+	dshash_seq_status			hstat;
+	TestDSMRegistryHashEntry   *entry;
+
+	if (flags & CHECKPOINT_END_OF_RECOVERY)
+		return;
+
+	tdr_attach_shmem();
+
+	file = AllocateFile(TDR_DUMP_FILE ".tmp", PG_BINARY_W);
+	if (file == NULL)
+		goto error;
+	if (fwrite(&TDR_FILE_HEADER, sizeof(uint32), 1, file) != 1)
+		goto error;
+
+	dshash_seq_init(&hstat, tdr_hash, false);
+	while ((entry = dshash_seq_next(&hstat)) != NULL)
+	{
+		int		keylen = offsetof(TestDSMRegistryHashEntry, val);
+		char   *val;
+		int32	vlen;
+
+		val = (char *) dsa_get_address(tdr_dsa, entry->val);
+		vlen = strlen(val) + 1;
+		if (fwrite(entry->key, keylen, 1, file) != 1 ||
+			fwrite(&vlen, sizeof(int32), 1, file) != 1 ||
+			fwrite(val, vlen, 1, file) != 1)
+		{
+			dshash_seq_term(&hstat);
+			goto error;
+		}
+	}
+	dshash_seq_term(&hstat);
+
+	if (FreeFile(file))
+	{
+		file = NULL;
+		goto error;
+	}
+
+	/*
+	 * Rename file into place, so we atomically replace any old one.
+	 */
+	(void) durable_rename(TDR_DUMP_FILE ".tmp", TDR_DUMP_FILE, LOG);
+	return;
+
+error:
+	ereport(LOG,
+			(errcode_for_file_access(),
+			 errmsg("could not write file \"%s\": %m",
+					TDR_DUMP_FILE ".tmp")));
+	if (file)
+		FreeFile(file);
+	unlink(TDR_DUMP_FILE ".tmp");
+}
+
+/*
+ * Entry point for this module.
+ */
+void
+_PG_init(void)
+{
+	prev_Checkpoint_hook = Checkpoint_hook;
+	Checkpoint_hook = pgss_Checkpoint;
+}
-- 
2.51.2