From 593d5a790fbfd54a984162b51ac75a81daff33f9 Mon Sep 17 00:00:00 2001
From: Ayush Tiwari <ayushtiwari.slg01@gmail.com>
Date: Sun, 9 Aug 2026 16:29:50 +0530
Subject: [PATCH v2 2/3] Roll back unfinished after-startup shmem
 initialization

An after-startup registration can fail after inserting some or all of its
areas into ShmemIndex but before init_fn completes.  Leaving those entries
visible lets a later backend attach to memory that was never fully initialized.
A partly inserted batch also cannot be retried because it contains a mixture
of existing and missing names.

On an error while creating new areas, remove all names in the request from
ShmemIndex and clear their backend-local handles.  ShmemIndexLock is still
held at this point, and none of the names existed before the attempt, so the
rollback cannot remove another subsystem's entry.  The allocated bytes cannot
be reclaimed, but they are no longer reachable through the index, and a later
attempt can reuse the names.

Do not roll back the attach path: those areas were initialized by another
backend and remain valid even if this backend's attach_fn fails.

Document this behavior and test partial allocation, init_fn failure, and
successful reuse of the rolled-back names.

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

diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index 2b8a11e7ad0..9269f72c847 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3742,6 +3742,13 @@ my_shmem_init(void *arg)
       lock (ShmemIndexLock), which prevents the race condition of two backends
       trying to initialize the memory area at the same time.
      </para>
+     <para>
+      If the call fails while creating the areas, their names are removed from
+      the shared memory index and the associated backend-local handles are
+      cleared, so that no backend can use memory that was not fully
+      initialized.  The allocation itself cannot be reclaimed, however, and
+      reduces the space available to a later attempt.
+     </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 a16e9388aad..ee619bd4e5a 100644
--- a/src/backend/storage/ipc/shmem.c
+++ b/src/backend/storage/ipc/shmem.c
@@ -213,6 +213,14 @@ enum shmem_request_state
 };
 static enum shmem_request_state shmem_request_state = SRS_INITIAL;
 
+/*
+ * True while an after-startup registration is creating its areas, that is,
+ * between the point where the first ShmemIndex entry is made and the point
+ * where the init_fn callback returns.  Until then the areas hold garbage, so
+ * an error in that window has to undo them.
+ */
+static bool late_shmem_init_in_progress = false;
+
 /*
  * This is the first data structure stored in the shared memory segment, at
  * the offset that PGShmemHeader->content_offset points to.  Allocations by
@@ -276,6 +284,7 @@ static bool firstNumaTouch = true;
 static void CallShmemCallbacksAfterStartup(const ShmemCallbacks *callbacks);
 static void CallShmemCallbacksAfterStartupInternal(const ShmemCallbacks *callbacks);
 static void DiscardPendingShmemRequests(void);
+static void RollbackLateShmemInit(void);
 static void InitShmemIndexEntry(ShmemRequest *request);
 static bool AttachShmemIndexEntry(ShmemRequest *request, bool missing_ok);
 
@@ -912,7 +921,10 @@ CallShmemCallbacksAfterStartup(const ShmemCallbacks *callbacks)
 	}
 	PG_FINALLY();
 	{
+		if (late_shmem_init_in_progress)
+			RollbackLateShmemInit();
 		DiscardPendingShmemRequests();
+		late_shmem_init_in_progress = false;
 		shmem_request_state = SRS_DONE;
 	}
 	PG_END_TRY();
@@ -924,7 +936,6 @@ CallShmemCallbacksAfterStartupInternal(const ShmemCallbacks *callbacks)
 	bool		found_any;
 	bool		notfound_any;
 
-
 	/*
 	 * Call the request callback first.  The callback makes ShmemRequest*()
 	 * calls for each shmem area, adding them to pending_shmem_requests.
@@ -963,6 +974,8 @@ CallShmemCallbacksAfterStartupInternal(const ShmemCallbacks *callbacks)
 	if (found_any && notfound_any)
 		elog(ERROR, "some of the requested shmem areas have already been initialized");
 
+	late_shmem_init_in_progress = !found_any;
+
 	/*
 	 * Allocate or attach all the shmem areas requested by the request_fn
 	 * callback.
@@ -973,7 +986,6 @@ CallShmemCallbacksAfterStartupInternal(const ShmemCallbacks *callbacks)
 			AttachShmemIndexEntry(request, false);
 		else
 			InitShmemIndexEntry(request);
-
 	}
 
 	/* Finish by calling the appropriate subsystem-specific callback */
@@ -988,9 +1000,56 @@ CallShmemCallbacksAfterStartupInternal(const ShmemCallbacks *callbacks)
 			callbacks->init_fn(callbacks->opaque_arg);
 	}
 
+	late_shmem_init_in_progress = false;
 	LWLockRelease(ShmemIndexLock);
 }
 
+/*
+ * Undo the visible parts of an initialization that did not finish.  The
+ * shared memory itself cannot be reclaimed, but removing the index entries
+ * prevents another backend from attaching to it and makes the names usable
+ * by a later attempt.
+ *
+ * This is called only on the create path, while ShmemIndexLock is still held,
+ * so any matching entries were inserted by this attempt.
+ */
+static void
+RollbackLateShmemInit(void)
+{
+	Assert(LWLockHeldByMe(ShmemIndexLock));
+
+	foreach_ptr(ShmemRequest, request, pending_shmem_requests)
+	{
+		if (hash_search(ShmemIndex, request->options->name,
+						HASH_REMOVE, NULL) == NULL)
+			continue;
+
+		switch (request->kind)
+		{
+			case SHMEM_KIND_STRUCT:
+				if (request->options->ptr)
+					*(request->options->ptr) = NULL;
+				break;
+			case SHMEM_KIND_HASH:
+				{
+					ShmemHashOpts *options = (ShmemHashOpts *) request->options;
+
+					if (options->ptr)
+						*(options->ptr) = NULL;
+					break;
+				}
+			case SHMEM_KIND_SLRU:
+				{
+					SlruOpts   *options = (SlruOpts *) request->options;
+
+					if (options->desc)
+						options->desc->shared = NULL;
+					break;
+				}
+		}
+	}
+}
+
 /* Release the requests accumulated by a request_fn callback. */
 static void
 DiscardPendingShmemRequests(void)
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 22c49686d73..a85d5df4d0f 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
@@ -69,6 +69,54 @@ 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");
 
+
+# A partly allocated batch is rolled back from ShmemIndex, although its bytes
+# cannot be reclaimed.  The same names can therefore be used by a smaller
+# retry, and the failed attempt leaves no usable local handle.
+my $result = $node->safe_psql("postgres", q[
+LOAD 'test_shmem';
+SET test_shmem.area_size = '1kB';
+SET test_shmem.extra_size = '1GB';
+DO $$ BEGIN PERFORM test_shmem_register(); EXCEPTION WHEN others THEN END $$;
+SELECT test_shmem_area_is_null(),
+			 (SELECT count(*)
+					FROM pg_shmem_allocations
+				 WHERE name IN ('test_shmem area 1024', 'test_shmem extra area'));]);
+is($result, "t|0", "partial allocation is hidden and its names are reusable");
+
+$result = $node->safe_psql("postgres", q[
+LOAD 'test_shmem';
+SET test_shmem.area_size = '1kB';
+DO $$ BEGIN PERFORM test_shmem_register(); END $$;
+SELECT get_test_shmem_attach_count();]);
+is($result, '0', "a smaller request can reuse a rolled-back name");
+
+SKIP:
+{
+	skip "injection points not supported by this build", 2
+	  if $ENV{enable_injection_points} ne 'yes';
+
+	$node->safe_psql("postgres", "CREATE EXTENSION IF NOT EXISTS injection_points;");
+	$result = $node->safe_psql("postgres", q[
+LOAD 'test_shmem';
+DO $$ BEGIN PERFORM injection_points_attach('test-shmem-init', 'error'); END $$;
+SET test_shmem.area_size = '2kB';
+DO $$ BEGIN PERFORM test_shmem_register(); EXCEPTION WHEN others THEN END $$;
+SELECT test_shmem_area_is_null(),
+			 (SELECT count(*)
+					FROM pg_shmem_allocations
+				 WHERE name = 'test_shmem area 2048');]);
+	is($result, "t|0", "an init callback failure is rolled back");
+	$node->safe_psql("postgres",
+		"SELECT injection_points_detach('test-shmem-init');");
+	$result = $node->safe_psql("postgres", q[
+LOAD 'test_shmem';
+SET test_shmem.area_size = '2kB';
+DO $$ BEGIN PERFORM test_shmem_register(); END $$;
+SELECT get_test_shmem_attach_count();]);
+	is($result, '0', "a name is reusable after an init callback failure");
+}
+
 # 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 27cb75f6020..7a358a2752b 100644
--- a/src/test/modules/test_shmem/test_shmem--1.0.sql
+++ b/src/test/modules/test_shmem/test_shmem--1.0.sql
@@ -11,3 +11,7 @@ AS 'MODULE_PATHNAME' LANGUAGE C;
 CREATE FUNCTION test_shmem_register()
 RETURNS pg_catalog.void STRICT
 AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION test_shmem_area_is_null()
+RETURNS pg_catalog.bool 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 c96e224a994..62cb7bc5645 100644
--- a/src/test/modules/test_shmem/test_shmem.c
+++ b/src/test/modules/test_shmem/test_shmem.c
@@ -40,6 +40,7 @@ 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 int	test_shmem_extra_size = 0;
 static char test_shmem_area_name[64] = "test_shmem area";
 
 static void test_shmem_request(void *arg);
@@ -67,6 +68,9 @@ test_shmem_request(void *arg)
 	ShmemRequestStruct(.name = test_shmem_area_name,
 					   .size = test_shmem_area_size,
 					   .ptr = (void **) &TestShmem);
+	if (test_shmem_extra_size > 0)
+		ShmemRequestStruct(.name = "test_shmem extra area",
+						   .size = test_shmem_extra_size);
 
 	if (test_shmem_after_startup)
 		INJECTION_POINT("test-shmem-request", NULL);
@@ -76,6 +80,8 @@ static void
 test_shmem_init(void *arg)
 {
 	elog(LOG, "init callback called");
+	if (test_shmem_after_startup)
+		INJECTION_POINT("test-shmem-init", NULL);
 	if (TestShmem->initialized)
 		elog(ERROR, "shmem area already initialized");
 	TestShmem->initialized = true;
@@ -112,6 +118,14 @@ _PG_init(void)
 							PGC_USERSET,
 							GUC_UNIT_BYTE,
 							NULL, NULL, NULL);
+	DefineCustomIntVariable("test_shmem.extra_size",
+							"Size of an additional area in the same request.",
+							NULL,
+							&test_shmem_extra_size,
+							0, 0, INT_MAX,
+							PGC_USERSET,
+							GUC_UNIT_BYTE,
+							NULL, NULL, NULL);
 	MarkGUCPrefixReserved("test_shmem");
 	RegisterShmemCallbacks(&TestShmemCallbacks);
 }
@@ -121,10 +135,18 @@ Datum
 test_shmem_register(PG_FUNCTION_ARGS)
 {
 	test_shmem_after_startup = true;
+	attached_or_initialized = false;
 	RegisterShmemCallbacks(&TestShmemCallbacks);
 	PG_RETURN_VOID();
 }
 
+PG_FUNCTION_INFO_V1(test_shmem_area_is_null);
+Datum
+test_shmem_area_is_null(PG_FUNCTION_ARGS)
+{
+	PG_RETURN_BOOL(TestShmem == NULL);
+}
+
 PG_FUNCTION_INFO_V1(get_test_shmem_attach_count);
 Datum
 get_test_shmem_attach_count(PG_FUNCTION_ARGS)
-- 
2.34.1

