From ddf6d1b38c01b36b28e98953e680aabf6aefbdbf Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakangas@iki.fi>
Date: Wed, 1 Apr 2026 18:10:31 +0300
Subject: [PATCH 03/14] Add test module to test after-startup shmem allocations

None of the existing modules could make use of the lazy shmem
allocation after postmaster startup:

- pg_stat_statements needs to load and dump stats file on startup and
  shutdown, which doesn't really work if the library is not loaded into
  postmaster

- test_aio registers injection points, which reference the library
  itself, which creates a weird initialization loop if you try to do
  that directly from _PG_init() in a backend. The initialization
  really needs to happen after _PG_init()

- injection_points would be a candidate, but it already knows to use
  DSM when it's not loaded from shared_preload_libraries.
---
 src/test/modules/Makefile                     |   1 +
 src/test/modules/meson.build                  |   1 +
 src/test/modules/test_shmem/Makefile          |  24 ++++
 src/test/modules/test_shmem/meson.build       |  33 ++++++
 .../test_shmem/t/001_late_shmem_alloc.pl      |  49 ++++++++
 .../modules/test_shmem/test_shmem--1.0.sql    |   9 ++
 src/test/modules/test_shmem/test_shmem.c      | 105 ++++++++++++++++++
 .../modules/test_shmem/test_shmem.control     |   3 +
 src/tools/pgindent/typedefs.list              |   1 +
 9 files changed, 226 insertions(+)
 create mode 100644 src/test/modules/test_shmem/Makefile
 create mode 100644 src/test/modules/test_shmem/meson.build
 create mode 100644 src/test/modules/test_shmem/t/001_late_shmem_alloc.pl
 create mode 100644 src/test/modules/test_shmem/test_shmem--1.0.sql
 create mode 100644 src/test/modules/test_shmem/test_shmem.c
 create mode 100644 src/test/modules/test_shmem/test_shmem.control

diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 864b407abcf..f1b04c99969 100644
--- 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 e5acacd5083..fc99552d9ab 100644
--- 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
index 00000000000..2407f7462fe
--- /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
index 00000000000..fb4bf328b8f
--- /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
index 00000000000..c154f57682a
--- /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
index 00000000000..2d01fd9256c
--- /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
index 00000000000..31c735be570
--- /dev/null
+++ b/src/test/modules/test_shmem/test_shmem.c
@@ -0,0 +1,105 @@
+/*-------------------------------------------------------------------------
+ *
+ * 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)
+{
+	static ShmemStructDesc TestShmemDesc;
+
+	elog(LOG, "test_shmem_request callback called");
+
+	ShmemRequestStruct(&TestShmemDesc,
+					   .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
index 00000000000..f2f26f4537a
--- /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 5894893997c..5358229b73f 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -3146,6 +3146,7 @@ TestDSMRegistryHashEntry
 TestDSMRegistryStruct
 TestDecodingData
 TestDecodingTxnData
+TestShmemData
 TestSpec
 TestValueType
 TextFreq
-- 
2.47.3

