From 76eb653e8bc317b4fdc47dc72748ec75a90c9361 Mon Sep 17 00:00:00 2001
From: Ayush Tiwari <ayushtiwari.slg01@gmail.com>
Date: Sun, 9 Aug 2026 16:23:02 +0530
Subject: [PATCH v2 1/3] Clean up pending shmem requests after an error

CallShmemCallbacksAfterStartup() cleared pending_shmem_requests and reset
shmem_request_state only on success.  Its request list lived in the caller's
memory context, so error cleanup could release the list while the global still
pointed at it.  A second registration attempt in the same backend then walked
a dangling list and failed the IsPointerList assertion.  The stale state also
made later RegisterShmemCallbacks() calls silently remember the callbacks for
later rather than run them.

Allocate the request records and list cells in TopMemoryContext, like their
options, and put the after-startup operation behind a PG_FINALLY block that
always frees the requests and restores the state.  Leave the startup paths
unchanged.

Add tests for errors in the request callback and during allocation.  They
retry in the same backend because the affected state is process-local.

Author: Ayush Tiwari <ayushtiwari.slg01@gmail.com>
---
 src/backend/storage/ipc/shmem.c               | 43 +++++++++++++---
 src/test/modules/test_shmem/Makefile          |  3 ++
 src/test/modules/test_shmem/meson.build       |  3 ++
 .../test_shmem/t/001_late_shmem_alloc.pl      | 51 +++++++++++++++++++
 .../modules/test_shmem/test_shmem--1.0.sql    |  4 ++
 src/test/modules/test_shmem/test_shmem.c      | 40 ++++++++++++++-
 6 files changed, 134 insertions(+), 10 deletions(-)

diff --git a/src/backend/storage/ipc/shmem.c b/src/backend/storage/ipc/shmem.c
index a3d56cf55dd..a16e9388aad 100644
--- a/src/backend/storage/ipc/shmem.c
+++ b/src/backend/storage/ipc/shmem.c
@@ -274,6 +274,8 @@ typedef struct
 static bool firstNumaTouch = true;
 
 static void CallShmemCallbacksAfterStartup(const ShmemCallbacks *callbacks);
+static void CallShmemCallbacksAfterStartupInternal(const ShmemCallbacks *callbacks);
+static void DiscardPendingShmemRequests(void);
 static void InitShmemIndexEntry(ShmemRequest *request);
 static bool AttachShmemIndexEntry(ShmemRequest *request, bool missing_ok);
 
@@ -336,6 +338,7 @@ void
 ShmemRequestInternal(ShmemStructOpts *options, ShmemRequestKind kind)
 {
 	ShmemRequest *request;
+	MemoryContext oldcontext;
 
 	/* Check the options */
 	if (options->name == NULL)
@@ -373,11 +376,13 @@ ShmemRequestInternal(ShmemStructOpts *options, ShmemRequestKind kind)
 							options->name)));
 	}
 
-	/* Request looks valid, remember it */
+	/* Keep the requests and list cells alive until we explicitly free them. */
+	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
 	request = palloc(sizeof(ShmemRequest));
 	request->options = options;
 	request->kind = kind;
 	pending_shmem_requests = lappend(pending_shmem_requests, request);
+	MemoryContextSwitchTo(oldcontext);
 }
 
 /*
@@ -896,12 +901,29 @@ RegisterShmemCallbacks(const ShmemCallbacks *callbacks)
  */
 static void
 CallShmemCallbacksAfterStartup(const ShmemCallbacks *callbacks)
+{
+	Assert(shmem_request_state == SRS_DONE);
+	Assert(pending_shmem_requests == NIL);
+
+	shmem_request_state = SRS_REQUESTING;
+	PG_TRY();
+	{
+		CallShmemCallbacksAfterStartupInternal(callbacks);
+	}
+	PG_FINALLY();
+	{
+		DiscardPendingShmemRequests();
+		shmem_request_state = SRS_DONE;
+	}
+	PG_END_TRY();
+}
+
+static void
+CallShmemCallbacksAfterStartupInternal(const ShmemCallbacks *callbacks)
 {
 	bool		found_any;
 	bool		notfound_any;
 
-	Assert(shmem_request_state == SRS_DONE);
-	shmem_request_state = SRS_REQUESTING;
 
 	/*
 	 * Call the request callback first.  The callback makes ShmemRequest*()
@@ -914,7 +936,6 @@ CallShmemCallbacksAfterStartup(const ShmemCallbacks *callbacks)
 
 	if (pending_shmem_requests == NIL)
 	{
-		shmem_request_state = SRS_DONE;
 		return;
 	}
 
@@ -953,10 +974,7 @@ CallShmemCallbacksAfterStartup(const ShmemCallbacks *callbacks)
 		else
 			InitShmemIndexEntry(request);
 
-		pfree(request->options);
 	}
-	list_free_deep(pending_shmem_requests);
-	pending_shmem_requests = NIL;
 
 	/* Finish by calling the appropriate subsystem-specific callback */
 	if (found_any)
@@ -971,7 +989,16 @@ CallShmemCallbacksAfterStartup(const ShmemCallbacks *callbacks)
 	}
 
 	LWLockRelease(ShmemIndexLock);
-	shmem_request_state = SRS_DONE;
+}
+
+/* Release the requests accumulated by a request_fn callback. */
+static void
+DiscardPendingShmemRequests(void)
+{
+	foreach_ptr(ShmemRequest, request, pending_shmem_requests)
+		pfree(request->options);
+	list_free_deep(pending_shmem_requests);
+	pending_shmem_requests = NIL;
 }
 
 /*
diff --git a/src/test/modules/test_shmem/Makefile b/src/test/modules/test_shmem/Makefile
index 2407f7462fe..fed8e29c8f5 100644
--- a/src/test/modules/test_shmem/Makefile
+++ b/src/test/modules/test_shmem/Makefile
@@ -2,6 +2,9 @@
 
 PGFILEDESC = "test_shmem - test code for shmem allocations"
 
+EXTRA_INSTALL = src/test/modules/injection_points
+export enable_injection_points
+
 MODULE_big = test_shmem
 OBJS = \
 	$(WIN32RES) \
diff --git a/src/test/modules/test_shmem/meson.build b/src/test/modules/test_shmem/meson.build
index fb4bf328b8f..8f98f2c4e31 100644
--- a/src/test/modules/test_shmem/meson.build
+++ b/src/test/modules/test_shmem/meson.build
@@ -26,6 +26,9 @@ tests += {
   'sd': meson.current_source_dir(),
   'bd': meson.current_build_dir(),
   'tap': {
+    'env': {
+      'enable_injection_points': get_option('injection_points') ? 'yes' : 'no',
+    },
     '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
index 5cf07d071ec..22c49686d73 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
@@ -18,6 +18,57 @@ $node->start;
 
 $node->safe_psql("postgres", "CREATE EXTENSION test_shmem;");
 
+# Run a registration twice in one backend.  The state under test is
+# backend-local, so separate safe_psql calls would not exercise a retry.
+sub register_twice
+{
+	my ($setup) = @_;
+	my $stderr;
+
+	$node->psql(
+		"postgres", qq[
+LOAD 'test_shmem';
+$setup
+DO \$\$
+BEGIN
+  FOR i IN 1..2 LOOP
+    BEGIN
+      PERFORM test_shmem_register();
+      RAISE NOTICE 'attempt %: ok', i;
+    EXCEPTION WHEN others THEN
+      RAISE NOTICE 'attempt %: %', i, SQLERRM;
+    END;
+  END LOOP;
+END
+\$\$;],
+		stderr => \$stderr);
+	return $stderr;
+}
+
+SKIP:
+{
+	skip "injection points not supported by this build", 2
+	  if $ENV{enable_injection_points} ne 'yes';
+
+	$node->safe_psql("postgres", "CREATE EXTENSION injection_points;");
+	my $stderr = register_twice(
+		"SELECT injection_points_attach('test-shmem-request', 'error');");
+	my @failures =
+	  ($stderr =~ /attempt \d: error triggered for injection point/g);
+	is(scalar @failures, 2,
+		"request callback failure is reported on both attempts");
+	unlike($stderr, qr/server closed the connection/,
+		"request callback failure does not crash the retry");
+	$node->safe_psql("postgres",
+		"SELECT injection_points_detach('test-shmem-request');");
+}
+
+my $stderr = register_twice("SET test_shmem.area_size = '1GB';");
+my @failures = ($stderr =~ /attempt \d: not enough shared memory/g);
+is(scalar @failures, 2, "allocation failure is reported on both attempts");
+unlike($stderr, qr/server closed the connection/,
+	"allocation failure does not crash the retry");
+
 # 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 2d01fd9256c..27cb75f6020 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()
+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 9bd4012b435..c96e224a994 100644
--- a/src/test/modules/test_shmem/test_shmem.c
+++ b/src/test/modules/test_shmem/test_shmem.c
@@ -17,9 +17,13 @@
 
 #include "postgres.h"
 
+#include <limits.h>
+
 #include "fmgr.h"
 #include "miscadmin.h"
 #include "storage/shmem.h"
+#include "utils/guc.h"
+#include "utils/injection_point.h"
 
 
 PG_MODULE_MAGIC;
@@ -34,6 +38,9 @@ typedef struct TestShmemData
 static TestShmemData *TestShmem;
 
 static bool attached_or_initialized = false;
+static int	test_shmem_area_size = sizeof(TestShmemData);
+static bool test_shmem_after_startup = false;
+static char test_shmem_area_name[64] = "test_shmem area";
 
 static void test_shmem_request(void *arg);
 static void test_shmem_init(void *arg);
@@ -51,9 +58,18 @@ test_shmem_request(void *arg)
 {
 	elog(LOG, "test_shmem_request callback called");
 
-	ShmemRequestStruct(.name = "test_shmem area",
-					   .size = sizeof(TestShmemData),
+	if (test_shmem_area_size == sizeof(TestShmemData))
+		strcpy(test_shmem_area_name, "test_shmem area");
+	else
+		snprintf(test_shmem_area_name, sizeof(test_shmem_area_name),
+				 "test_shmem area %d", test_shmem_area_size);
+
+	ShmemRequestStruct(.name = test_shmem_area_name,
+					   .size = test_shmem_area_size,
 					   .ptr = (void **) &TestShmem);
+
+	if (test_shmem_after_startup)
+		INJECTION_POINT("test-shmem-request", NULL);
 }
 
 static void
@@ -86,7 +102,27 @@ void
 _PG_init(void)
 {
 	elog(LOG, "test_shmem module's _PG_init called");
+
+	DefineCustomIntVariable("test_shmem.area_size",
+							"Size of the shmem area to request.",
+							NULL,
+							&test_shmem_area_size,
+							sizeof(TestShmemData),
+							sizeof(TestShmemData), INT_MAX,
+							PGC_USERSET,
+							GUC_UNIT_BYTE,
+							NULL, NULL, NULL);
+	MarkGUCPrefixReserved("test_shmem");
+	RegisterShmemCallbacks(&TestShmemCallbacks);
+}
+
+PG_FUNCTION_INFO_V1(test_shmem_register);
+Datum
+test_shmem_register(PG_FUNCTION_ARGS)
+{
+	test_shmem_after_startup = true;
 	RegisterShmemCallbacks(&TestShmemCallbacks);
+	PG_RETURN_VOID();
 }
 
 PG_FUNCTION_INFO_V1(get_test_shmem_attach_count);
-- 
2.34.1

