Add a test module to test after-startup shmem allocations

The old ShmemInit{Struct/Hash}() functions could be used after
postmaster statup, as long as the allocation is small enough to fit in
spare shmem reserved at startup. I believe some extensions do that,
although we hadn't really documented it and had not coverage for it.
The new test module covers that after-startup usage with the new
ShmemRequestStruct() functions.

Reviewed-by: Ashutosh Bapat <ashutosh.bapat.oss@gmail.com>
Reviewed-by: Matthias van de Meent <boekewurm+postgres@gmail.com>
Reviewed-by: Daniel Gustafsson <daniel@yesql.se>
Discussion: https://www.postgresql.org/message-id/CAExHW5vM1bneLYfg0wGeAa=52UiJ3z4vKd3AJ72X8Fw6k3KKrg@mail.gmail.com

diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 864b407abcff701896bc8cd19c2fe9e4ef1393e3..f1b04c99969c429683976a72416dbe6b455fb67b 100644 (file)
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -48,6 +48,7 @@ SUBDIRS = \
                  test_resowner \
                  test_rls_hooks \
                  test_saslprep \
+                 test_shmem \
                  test_shm_mq \
                  test_slru \
                  test_tidstore \

diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build
index e5acacd508368d596a181df74d0226551f63ca9a..fc99552d9ab9af7bf2f44d18c0d8e407241fb2ea 100644 (file)
--- a/src/test/modules/meson.build
+++ b/src/test/modules/meson.build
@@ -49,6 +49,7 @@ subdir('test_regex')
 subdir('test_resowner')
 subdir('test_rls_hooks')
 subdir('test_saslprep')
+subdir('test_shmem')
 subdir('test_shm_mq')
 subdir('test_slru')
 subdir('test_tidstore')

diff --git a/src/test/modules/test_shmem/Makefile b/src/test/modules/test_shmem/Makefile
new file mode 100644 (file)
index 0000000..2407f74
--- /dev/null
+++ b/src/test/modules/test_shmem/Makefile
@@ -0,0 +1,24 @@
+# src/test/modules/test_shmem/Makefile
+
+PGFILEDESC = "test_shmem - test code for shmem allocations"
+
+MODULE_big = test_shmem
+OBJS = \
+       $(WIN32RES) \
+       test_shmem.o
+
+EXTENSION = test_shmem
+DATA = test_shmem--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_shmem
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif

diff --git a/src/test/modules/test_shmem/meson.build b/src/test/modules/test_shmem/meson.build
new file mode 100644 (file)
index 0000000..fb4bf32
--- /dev/null
+++ b/src/test/modules/test_shmem/meson.build
@@ -0,0 +1,33 @@
+# Copyright (c) 2024-2026, PostgreSQL Global Development Group
+
+test_shmem_sources = files(
+  'test_shmem.c',
+)
+
+if host_system == 'windows'
+  test_shmem_sources += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_shmem',
+    '--FILEDESC', 'test_shmem - test code for shmem allocations',])
+endif
+
+test_shmem = shared_module('test_shmem',
+  test_shmem_sources,
+  kwargs: pg_test_mod_args,
+)
+test_install_libs += test_shmem
+
+test_install_data += files(
+  'test_shmem.control',
+  'test_shmem--1.0.sql',
+)
+
+tests += {
+  'name': 'test_shmem',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'tap': {
+    'tests': [
+      't/001_late_shmem_alloc.pl',
+    ],
+  },
+}

diff --git a/src/test/modules/test_shmem/t/001_late_shmem_alloc.pl b/src/test/modules/test_shmem/t/001_late_shmem_alloc.pl
new file mode 100644 (file)
index 0000000..c154f57
--- /dev/null
+++ b/src/test/modules/test_shmem/t/001_late_shmem_alloc.pl
@@ -0,0 +1,49 @@
+# Copyright (c) 2025-2026, PostgreSQL Global Development Group
+
+use strict;
+use warnings FATAL => 'all';
+
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+###
+# Test allocating memory after startup, i.e. when the library is not
+# in shared_preload_libraries
+###
+my $node = PostgreSQL::Test::Cluster->new('main');
+$node->init;
+$node->start;
+
+
+$node->safe_psql("postgres", "CREATE EXTENSION test_shmem;");
+
+# Check that the attach counter is incremented on a new connection
+my $attach_count1 = $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();");
+my $attach_count2 = $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();");
+cmp_ok($attach_count2, '>', $attach_count1, "attach callback is called in each backend");
+$node->stop;
+
+###
+# Test that loading via shared_preload_libraries also works
+###
+$node->append_conf('postgresql.conf', "shared_preload_libraries = 'test_shmem'");
+$node->start;
+
+# When loaded via shared_preload_libraries, the attach callback is
+# called or not, depending on whether this is an EXEC_BACKEND build.
+my $exec_backend = $node->safe_psql("postgres", "SHOW debug_exec_backend;") eq 'on';
+$attach_count1 = $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();");
+$attach_count2 = $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();");
+
+if ($exec_backend)
+{
+   cmp_ok($attach_count2, '>', $attach_count1, "attach callback is called in each backend when loaded via shared_preload_libraries");
+}
+else
+{
+   ok($attach_count1 == 0 && $attach_count2 == 0, "attach callback is not called when loaded via shared_preload_libraries");
+}
+
+$node->stop;
+done_testing();

diff --git a/src/test/modules/test_shmem/test_shmem--1.0.sql b/src/test/modules/test_shmem/test_shmem--1.0.sql
new file mode 100644 (file)
index 0000000..2d01fd9
--- /dev/null
+++ b/src/test/modules/test_shmem/test_shmem--1.0.sql
@@ -0,0 +1,9 @@
+/* src/test/modules/test_shmem/test_shmem--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_shmem" to load this file. \quit
+
+
+CREATE FUNCTION get_test_shmem_attach_count()
+RETURNS pg_catalog.int4 STRICT
+AS 'MODULE_PATHNAME' LANGUAGE C;

diff --git a/src/test/modules/test_shmem/test_shmem.c b/src/test/modules/test_shmem/test_shmem.c
new file mode 100644 (file)
index 0000000..9bd4012
--- /dev/null
+++ b/src/test/modules/test_shmem/test_shmem.c
@@ -0,0 +1,101 @@
+/*-------------------------------------------------------------------------
+ *
+ * test_shmem.c
+ *             Helpers to test shmem allocation routines
+ *
+ * Test basic memory allocation in an extension module. One notable feature
+ * that is not exercised by any other module in the repository is the
+ * allocating (non-DSM) shared memory after postmaster startup.
+ *
+ * Copyright (c) 2020-2026, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *       src/test/modules/test_shmem/test_shmem.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "storage/shmem.h"
+
+
+PG_MODULE_MAGIC;
+
+typedef struct TestShmemData
+{
+       int                     value;
+       bool            initialized;
+       int                     attach_count;
+} TestShmemData;
+
+static TestShmemData *TestShmem;
+
+static bool attached_or_initialized = false;
+
+static void test_shmem_request(void *arg);
+static void test_shmem_init(void *arg);
+static void test_shmem_attach(void *arg);
+
+static const ShmemCallbacks TestShmemCallbacks = {
+       .flags = SHMEM_CALLBACKS_ALLOW_AFTER_STARTUP,
+       .request_fn = test_shmem_request,
+       .init_fn = test_shmem_init,
+       .attach_fn = test_shmem_attach,
+};
+
+static void
+test_shmem_request(void *arg)
+{
+       elog(LOG, "test_shmem_request callback called");
+
+       ShmemRequestStruct(.name = "test_shmem area",
+                                          .size = sizeof(TestShmemData),
+                                          .ptr = (void **) &TestShmem);
+}
+
+static void
+test_shmem_init(void *arg)
+{
+       elog(LOG, "init callback called");
+       if (TestShmem->initialized)
+               elog(ERROR, "shmem area already initialized");
+       TestShmem->initialized = true;
+
+       if (attached_or_initialized)
+               elog(ERROR, "attach or initialize already called in this process");
+       attached_or_initialized = true;
+}
+
+static void
+test_shmem_attach(void *arg)
+{
+       elog(LOG, "test_shmem_attach callback called");
+       if (!TestShmem->initialized)
+               elog(ERROR, "shmem area not yet initialized");
+       TestShmem->attach_count++;
+
+       if (attached_or_initialized)
+               elog(ERROR, "attach or initialize already called in this process");
+       attached_or_initialized = true;
+}
+
+void
+_PG_init(void)
+{
+       elog(LOG, "test_shmem module's _PG_init called");
+       RegisterShmemCallbacks(&TestShmemCallbacks);
+}
+
+PG_FUNCTION_INFO_V1(get_test_shmem_attach_count);
+Datum
+get_test_shmem_attach_count(PG_FUNCTION_ARGS)
+{
+       if (!attached_or_initialized)
+               elog(ERROR, "shmem area not attached or initialized in this process");
+       if (!TestShmem->initialized)
+               elog(ERROR, "shmem area not yet initialized");
+       PG_RETURN_INT32(TestShmem->attach_count);
+}

diff --git a/src/test/modules/test_shmem/test_shmem.control b/src/test/modules/test_shmem/test_shmem.control
new file mode 100644 (file)
index 0000000..f2f26f4
--- /dev/null
+++ b/src/test/modules/test_shmem/test_shmem.control
@@ -0,0 +1,3 @@
+comment = 'Test code for shmem allocations'
+default_version = '1.0'
+module_pathname = '$libdir/test_shmem'

diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 8eea59b62fca59407d7b2a0f41f0e695d3213b7b..298d6ac1639e71365635f2560a904388e437485a 100644 (file)
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -3147,6 +3147,7 @@ TestDSMRegistryHashEntry
 TestDSMRegistryStruct
 TestDecodingData
 TestDecodingTxnData
+TestShmemData
 TestSpec
 TestValueType
 TextFreq