From 904e010acc467ba291f1d29248374c7f917fe090 Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakangas@iki.fi>
Date: Thu, 24 Sep 2026 16:20:51 +0300
Subject: [PATCH v2 1/1] Allow unknown-size shmem attachments in single-user
 mode

ShmemRequestInternal() rejected SHMEM_ATTACH_UNKNOWN_SIZE whenever
IsUnderPostmaster was false.  That also rejected an after-startup
attachment in a standalone backend, even if the area exists.

To fix, check whether the shared memory has been initialized
(ShmemIndex == NULL) rather than IsUnderPostmaster.

Refactor the tests, and add more test cases for after-startup
attachment, with SHMEM_ATTACH_UNKNOWN_SIZE and a mismatching size, and
for this single-user case.

Author: Ayush Tiwari <ayushtiwari.slg01@gmail.com>
Discusssion: https://www.postgresql.org/message-id/CAJTYsWXREGXCnWdZJsXt8pT2UY_Y2GSx+HOUu2X3s8YDCLj0_A@mail.gmail.com
Backpatch-through: 19
---
 src/backend/storage/ipc/shmem.c               |  24 ++-
 .../test_shmem/t/001_late_shmem_alloc.pl      | 144 ++++++++++--------
 .../modules/test_shmem/test_shmem--1.0.sql    |   4 +
 src/test/modules/test_shmem/test_shmem.c      |  55 +++++++
 4 files changed, 152 insertions(+), 75 deletions(-)

diff --git a/src/backend/storage/ipc/shmem.c b/src/backend/storage/ipc/shmem.c
index 0e827344462..9ceff1723ed 100644
--- a/src/backend/storage/ipc/shmem.c
+++ b/src/backend/storage/ipc/shmem.c
@@ -348,33 +348,27 @@ ShmemRequestInternal(ShmemStructOpts *options, ShmemRequestKind kind)
 	MemoryContext oldcontext;
 	ShmemRequest *request;
 
+	/* Check that we're in the right state */
+	if (shmem_request_state != SRS_REQUESTING)
+		elog(ERROR, "ShmemRequestStruct can only be called from a shmem_request callback");
+
 	/* Check the options */
 	if (options->name == NULL)
 		elog(ERROR, "shared memory request is missing 'name' option");
 
-	if (IsUnderPostmaster)
-	{
-		if (options->size <= 0 && options->size != SHMEM_ATTACH_UNKNOWN_SIZE)
-			elog(ERROR, "invalid size %zd for shared memory request for \"%s\"",
-				 options->size, options->name);
-	}
-	else
+	if (options->size == SHMEM_ATTACH_UNKNOWN_SIZE)
 	{
-		if (options->size == SHMEM_ATTACH_UNKNOWN_SIZE)
+		if (ShmemIndex == NULL)
 			elog(ERROR, "SHMEM_ATTACH_UNKNOWN_SIZE cannot be used during startup");
-		if (options->size <= 0)
-			elog(ERROR, "invalid size %zd for shared memory request for \"%s\"",
-				 options->size, options->name);
 	}
+	else if (options->size <= 0)
+		elog(ERROR, "invalid size %zd for shared memory request for \"%s\"",
+			 options->size, options->name);
 
 	if (options->alignment != 0 && pg_nextpower2_size_t(options->alignment) != options->alignment)
 		elog(ERROR, "invalid alignment %zu for shared memory request for \"%s\"",
 			 options->alignment, options->name);
 
-	/* Check that we're in the right state */
-	if (shmem_request_state != SRS_REQUESTING)
-		elog(ERROR, "ShmemRequestStruct can only be called from a shmem_request callback");
-
 	/* Check that it's not already registered in this process */
 	foreach_ptr(ShmemRequest, existing, pending_shmem_requests)
 	{
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 7d961fe2790..c52bbed7bf5 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
@@ -16,62 +16,7 @@ $node->safe_psql("postgres", "CREATE EXTENSION test_shmem");
 $node->stop;
 
 ###
-# Test allocating memory after startup, i.e. when the library is not
-# in shared_preload_libraries
-###
-$node->start;
-
-# 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 trying to allocate a new shmem area with size =
-# SHMEM_ATTACH_UNKNOWN_SIZE (-1) fails.
-###
-$node->append_conf('postgresql.conf', "test_shmem.area_size = -1");
-$node->start;
-
-my (undef, undef, $stderr) =
-  $node->psql("postgres", "SELECT get_test_shmem_attach_count();");
-like(
-	$stderr,
-	qr/cannot attach to shared memory struct "test_shmem area" because it does not exist/,
-	"unknown size request for a nonexistent area fails");
-
-$node->stop;
-$node->adjust_conf('postgresql.conf', 'test_shmem.area_size', undef);
-
-###
-# Test allocating memory after startup in single-user mode
-###
-SKIP:
-{
-	# Skip the test on Windows, as single-user mode would fail on permission
-	# failure with privileged accounts.
-	skip 'single-user test is not supported by this platform', 1
-	  if $windows_os;
-	my $query = "SELECT get_test_shmem_attach_count();\n";
-	my $result = run_log(
-		[
-			'postgres', '--single', '-F',
-			'-c' => 'exit_on_error=true',
-			'-D' => $node->data_dir,
-			'postgres'
-		],
-		'<' => \$query);
-
-	ok($result, "shmem area is initialized in single-user mode");
-}
-
-###
-# Test that loading via shared_preload_libraries also works
+# Test that loading via shared_preload_libraries works
 ###
 $node->append_conf('postgresql.conf',
 	"shared_preload_libraries = 'test_shmem'");
@@ -81,9 +26,9 @@ $node->start;
 # 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 =
+my $attach_count1 =
   $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();");
-$attach_count2 =
+my $attach_count2 =
   $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();");
 
 if ($exec_backend)
@@ -99,9 +44,8 @@ else
 	);
 }
 
-# clean up
 $node->stop;
-$node->adjust_conf('postgresql.conf', "shared_preload_libraries", undef);
+$node->adjust_conf('postgresql.conf', 'shared_preload_libraries', undef);
 
 ###
 # Test a failure in initializing the shared memory area
@@ -139,6 +83,49 @@ SKIP:
 	$node->stop;
 }
 
+###
+# Test allocating memory after startup, i.e. when the library is not
+# in shared_preload_libraries
+###
+$node->start;
+
+# This first call to the function after startup loads the library
+# and initializes the shmem area.
+$attach_count1 =
+  $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();");
+
+# Check that the attach counter is incremented on a new connection
+$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");
+
+# Allocate another shmem area, after the library is loaded.
+my $stderr;
+my $res = $node->safe_psql("postgres",
+	"SELECT test_shmem_register('test_shmem after startup', 20, 1);");
+is($res, 0, 'allocate after startup');
+
+# Test attaching to it again
+$res = $node->safe_psql("postgres",
+	"SELECT test_shmem_register('test_shmem after startup', 20, 2);");
+is($res, 1, 'attach after startup');
+
+# If the size doesn't match when attaching, you get an error
+(undef, undef, $stderr) =
+  $node->psql("postgres", "SELECT test_shmem_register('test_shmem after startup', 25, 3);");
+like(
+	$stderr,
+	qr/ERROR:  shared memory struct "test_shmem after startup" was created with different size: existing 20, requested 25/,
+	"attaching with different size fails");
+
+# Test attaching with SHMEM_ATTACH_UNKNOWN_SIZE
+$res =
+  $node->safe_psql("postgres", "SELECT test_shmem_register('test_shmem after startup', -1, 4);");
+is($res, 2, 'attach with SHMEM_ATTACH_UNKNOWN_SIZE');
+
+$node->stop;
+
 ###
 # Test "out of shared memory" in an after-startup request
 ###
@@ -167,4 +154,41 @@ $session->query_safe("SELECT get_test_shmem_attach_count();");
 $session->quit;
 $node->stop;
 
+
+###
+# Test allocating memory after startup in single-user mode
+###
+SKIP:
+{
+	# Skip the test on Windows, as single-user mode would fail on permission
+	# failure with privileged accounts.
+	skip 'single-user test is not supported by this platform', 1
+	  if $windows_os;
+
+	my @command = (
+		'postgres', '--single', '-F',
+		'-c' => 'exit_on_error=true',
+		'-D' => $node->data_dir,
+		'postgres');
+
+	my $queries = "SELECT get_test_shmem_attach_count();\n";
+	my $result = run_log([@command], '<' => \$queries);
+	ok($result, "shmem area is initialized in single-user mode");
+
+	$queries = qq{
+-- allocate
+SELECT test_shmem_register('test_shmem after startup', 25, 1);
+-- attach
+SELECT test_shmem_register('test_shmem after startup', 25, 2);
+-- attach with SHMEM_ATTACH_UNKNOWN_SIZE
+SELECT test_shmem_register('test_shmem after startup', -1, 3);
+};
+	$result = run_log([@command], '<' => \$queries);
+	ok($result, "shmem area is initialized in single-user mode");
+}
+
+# clean up
+$node->stop;
+$node->adjust_conf('postgresql.conf', "shared_preload_libraries", undef);
+
 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
index 2d01fd9256c..1dec3ce0d68 100644
--- a/src/test/modules/test_shmem/test_shmem--1.0.sql
+++ b/src/test/modules/test_shmem/test_shmem--1.0.sql
@@ -7,3 +7,7 @@
 CREATE FUNCTION get_test_shmem_attach_count()
 RETURNS pg_catalog.int4 STRICT
 AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION test_shmem_register(pg_catalog.text, pg_catalog.int8, pg_catalog.int4)
+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
index 2a7e13512bf..772661abce7 100644
--- a/src/test/modules/test_shmem/test_shmem.c
+++ b/src/test/modules/test_shmem/test_shmem.c
@@ -20,6 +20,7 @@
 #include "fmgr.h"
 #include "miscadmin.h"
 #include "storage/shmem.h"
+#include "utils/builtins.h"
 #include "utils/guc.h"
 #include "utils/injection_point.h"
 
@@ -129,3 +130,57 @@ get_test_shmem_attach_count(PG_FUNCTION_ARGS)
 		elog(ERROR, "shmem area not yet initialized");
 	PG_RETURN_INT32(TestShmem->attach_count);
 }
+
+
+/*
+ * Callback for test_shmem_register().  test_shmem_register() provides the
+ * options, we just pass them through to ShmemRequestStruct.
+ */
+static void
+test_shmem_after_startup_request(void *arg)
+{
+	ShmemStructOpts *opts = (ShmemStructOpts *) arg;
+
+	elog(LOG, "test_shmem_after_startup_request callback called");
+
+	ShmemRequestStructWithOpts(opts);
+}
+
+/*
+ * Allocate or attach to a shmem segment, with the caller-supplied name and
+ * size.
+ *
+ * The given integer 'new_value' is stored in the segment, and the old value
+ * is returned.
+ */
+PG_FUNCTION_INFO_V1(test_shmem_register);
+Datum
+test_shmem_register(PG_FUNCTION_ARGS)
+{
+	char	   *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	int64		size = PG_GETARG_INT64(1);
+	int			new_value = PG_GETARG_INT32(2);
+	int			old_value;
+	int		   *attached = NULL;
+
+	ShmemStructOpts opts = {
+		.name = name,
+		.size = size,
+		.ptr = (void **) &attached,
+	};
+
+	ShmemCallbacks callbacks = {
+		.flags = SHMEM_CALLBACKS_ALLOW_AFTER_STARTUP,
+		.request_fn = test_shmem_after_startup_request,
+		.opaque_arg = &opts,
+	};
+
+	RegisterShmemCallbacks(&callbacks);
+	if (attached == NULL)
+		elog(ERROR, "could not attach to shared memory");
+
+	old_value = *attached;
+	*attached = new_value;
+
+	PG_RETURN_INT32(old_value);
+}
-- 
2.47.3

