From f92abbcc3667103628608d248870867200087e16 Mon Sep 17 00:00:00 2001 From: "Andrei V. Lepikhov" 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