From e3795b162dc8a316966b5e197828b3df16ecf072 Mon Sep 17 00:00:00 2001
From: Ayush Tiwari <ayushtiwari.slg01@gmail.com>
Date: Sun, 9 Aug 2026 16:34:27 +0530
Subject: [PATCH v2 3/3] Reject legacy shmem allocation from init and attach
 callbacks

After startup, init_fn and attach_fn run while ShmemIndexLock is held.
ShmemInitStruct() takes the same lock, so calling it from either callback
asserts in an assert-enabled build and deadlocks otherwise.

Reject ShmemInitStruct() in the callback states with a direct error.  This
also covers ShmemInitHash(), which calls ShmemInitStruct().  Legacy calls
during postmaster startup remain allowed for compatibility.

Document the restriction and test the init_fn case.

Author: Ayush Tiwari <ayushtiwari.slg01@gmail.com>
---
 doc/src/sgml/xfunc.sgml                               |  7 +++++++
 src/backend/storage/ipc/shmem.c                       |  4 ++++
 src/test/modules/test_shmem/t/001_late_shmem_alloc.pl | 11 +++++++++++
 src/test/modules/test_shmem/test_shmem--1.0.sql       |  2 +-
 src/test/modules/test_shmem/test_shmem.c              |  8 ++++++++
 5 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index 9269f72c847..a3dc603049d 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3749,6 +3749,13 @@ my_shmem_init(void *arg)
       initialized.  The allocation itself cannot be reclaimed, however, and
       reduces the space available to a later attempt.
      </para>
+     <para>
+      When the callbacks are run after startup, the legacy
+      <function>ShmemInitStruct()</function> and
+      <function>ShmemInitHash()</function> functions cannot be called from
+      <function>init_fn</function> or <function>attach_fn</function>.  Request
+      all areas from <function>request_fn</function> instead.
+     </para>
     </sect3>
 
     <sect3 id="xfunc-shared-addin-dynamic">
diff --git a/src/backend/storage/ipc/shmem.c b/src/backend/storage/ipc/shmem.c
index ee619bd4e5a..bed92e6d62f 100644
--- a/src/backend/storage/ipc/shmem.c
+++ b/src/backend/storage/ipc/shmem.c
@@ -1106,6 +1106,10 @@ ShmemInitStruct(const char *name, Size size, bool *foundPtr)
 	};
 	ShmemRequest request = {&options, SHMEM_KIND_STRUCT};
 
+	if (shmem_request_state == SRS_ATTACHING ||
+		shmem_request_state == SRS_AFTER_STARTUP_ATTACH_OR_INIT)
+		elog(ERROR, "cannot call ShmemInitStruct() from a shmem init or attach callback");
+
 	Assert(shmem_request_state == SRS_DONE ||
 		   shmem_request_state == SRS_INITIALIZING ||
 		   shmem_request_state == SRS_REQUESTING);
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
index a85d5df4d0f..bbbec4dda92 100644
--- a/src/test/modules/test_shmem/t/001_late_shmem_alloc.pl
+++ b/src/test/modules/test_shmem/t/001_late_shmem_alloc.pl
@@ -117,6 +117,17 @@ SELECT get_test_shmem_attach_count();]);
 	is($result, '0', "a name is reusable after an init callback failure");
 }
 
+
+my (undef, undef, $legacy_stderr) = $node->psql(
+	"postgres", q[
+LOAD 'test_shmem';
+SET test_shmem.area_size = '3kB';
+SELECT test_shmem_register(true);],
+	timeout => $PostgreSQL::Test::Utils::timeout_default);
+like($legacy_stderr,
+	qr/cannot call ShmemInitStruct\(\) from a shmem init or attach callback/,
+	"legacy allocation from an init callback is rejected");
+
 # Check that the attach counter is incremented on a new connection
 my $attach_count1 =
   $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();");
diff --git a/src/test/modules/test_shmem/test_shmem--1.0.sql b/src/test/modules/test_shmem/test_shmem--1.0.sql
index 7a358a2752b..ca42274d7bc 100644
--- a/src/test/modules/test_shmem/test_shmem--1.0.sql
+++ b/src/test/modules/test_shmem/test_shmem--1.0.sql
@@ -8,7 +8,7 @@ CREATE FUNCTION get_test_shmem_attach_count()
 RETURNS pg_catalog.int4 STRICT
 AS 'MODULE_PATHNAME' LANGUAGE C;
 
-CREATE FUNCTION test_shmem_register()
+CREATE FUNCTION test_shmem_register(pg_catalog.bool DEFAULT false)
 RETURNS pg_catalog.void 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
index 62cb7bc5645..1322c457e34 100644
--- a/src/test/modules/test_shmem/test_shmem.c
+++ b/src/test/modules/test_shmem/test_shmem.c
@@ -41,6 +41,7 @@ static bool attached_or_initialized = false;
 static int	test_shmem_area_size = sizeof(TestShmemData);
 static bool test_shmem_after_startup = false;
 static int	test_shmem_extra_size = 0;
+static bool test_shmem_legacy_init = false;
 static char test_shmem_area_name[64] = "test_shmem area";
 
 static void test_shmem_request(void *arg);
@@ -82,6 +83,12 @@ test_shmem_init(void *arg)
 	elog(LOG, "init callback called");
 	if (test_shmem_after_startup)
 		INJECTION_POINT("test-shmem-init", NULL);
+	if (test_shmem_legacy_init)
+	{
+		bool		found;
+
+		(void) ShmemInitStruct("test_shmem legacy area", 1024, &found);
+	}
 	if (TestShmem->initialized)
 		elog(ERROR, "shmem area already initialized");
 	TestShmem->initialized = true;
@@ -134,6 +141,7 @@ PG_FUNCTION_INFO_V1(test_shmem_register);
 Datum
 test_shmem_register(PG_FUNCTION_ARGS)
 {
+	test_shmem_legacy_init = PG_NARGS() > 0 && PG_GETARG_BOOL(0);
 	test_shmem_after_startup = true;
 	attached_or_initialized = false;
 	RegisterShmemCallbacks(&TestShmemCallbacks);
-- 
2.34.1

