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

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 \

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')

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

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',
+    ],
+  },
+}

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();

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;

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);
+}

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'

src/tools/pgindent/typedefs.list

@@ -3147,6 +3147,7 @@ TestDSMRegistryHashEntry
 TestDSMRegistryStruct
 TestDecodingData
 TestDecodingTxnData
+TestShmemData
 TestSpec
 TestValueType
 TextFreq