From 78e61dac82d2143527228b4bb9495989e4e1d497 Mon Sep 17 00:00:00 2001
From: Ayush Tiwari <ayushtiwari.slg01@gmail.com>
Date: Thu, 6 Aug 2026 13:33:00 +0530
Subject: [PATCH 2/3] Preflight after-startup shmem batches for available space

An after-startup request can contain several shared memory areas.  They
were allocated one at a time, so if the remaining shared memory ran out
on a later area, earlier ones remained allocated and registered.  A
retry then found some names present and some missing and failed with
"some of the requested shmem areas have already been initialized".
Since shared memory cannot be freed, the subsystem stayed wedged until
server restart.

Check the remaining shared memory against the whole batch before
creating any area.  Named allocations are serialized by ShmemIndexLock;
low-level direct ShmemAlloc() callers can still race the preflight.

Add a test with one small area followed by an oversized one, and verify
that the preflight creates no entries and can be retried.

Author: Ayush Tiwari <ayushtiwari.slg01@gmail.com>
---
 src/backend/storage/ipc/shmem.c               | 32 +++++++++++++++++++
 .../test_shmem/t/001_late_shmem_alloc.pl      | 13 ++++++++
 src/test/modules/test_shmem/test_shmem.c      |  7 ++++
 3 files changed, 52 insertions(+)

diff --git a/src/backend/storage/ipc/shmem.c b/src/backend/storage/ipc/shmem.c
index 9986d6a69ce..c117955e49b 100644
--- a/src/backend/storage/ipc/shmem.c
+++ b/src/backend/storage/ipc/shmem.c
@@ -991,6 +991,38 @@ CallShmemCallbacksAfterStartup(const ShmemCallbacks *callbacks)
 	if (found_any && notfound_any)
 		elog(ERROR, "some of the requested shmem areas have already been initialized");
 
+	/*
+	 * Check that the whole batch fits before creating any of it.  Otherwise a
+	 * failure partway through leaves a mix of present and missing areas, which
+	 * cannot be retried.  Low-level ShmemAlloc() callers can still race this
+	 * check, but named allocations are serialized by ShmemIndexLock.
+	 */
+	if (!found_any)
+	{
+		Size		offset;
+
+		SpinLockAcquire(&ShmemAllocator->shmem_lock);
+		offset = ShmemAllocator->free_offset;
+		SpinLockRelease(&ShmemAllocator->shmem_lock);
+
+		foreach_ptr(ShmemRequest, request, pending_shmem_requests)
+		{
+			Size		alignment = request->options->alignment;
+
+			if (alignment < PG_CACHE_LINE_SIZE)
+				alignment = PG_CACHE_LINE_SIZE;
+			offset = TYPEALIGN(alignment, offset);
+			offset = add_size(offset, request->options->size);
+			if (offset > ShmemSegHdr->totalsize)
+				ereport(ERROR,
+						(errcode(ERRCODE_OUT_OF_MEMORY),
+						 errmsg("not enough shared memory for data structure"
+								" \"%s\" (%zd bytes requested)",
+								request->options->name,
+								request->options->size)));
+		}
+	}
+
 	/*
 	 * Allocate or attach all the shmem areas requested by the request_fn
 	 * callback.
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 a9481997d52..5915f060a0c 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
@@ -53,6 +53,19 @@ foreach my $mode (0, 1)
 	like($stderr, qr/attempt 1: /, "shmem request $mode fails");
 	like($stderr, qr/attempt 2: /, "shmem request $mode fails when retried");
 }
+
+my ($ret, $stdout, $stderr) = try_shmem_failure_twice(2);
+is($ret, 0, 'session survives a partly oversized request batch');
+like($stderr, qr/attempt 2: .*not enough shared memory/,
+	'a partly oversized batch can be retried');
+unlike($stderr, qr/already been initialized/,
+	'a partly oversized batch does not wedge later attempts');
+is( $node->safe_psql(
+		'postgres',
+		"SELECT count(*) FROM pg_shmem_allocations WHERE name LIKE 'test_shmem partial%';"
+	),
+	'0',
+	'a partly oversized batch creates no areas');
 $node->stop;
 
 ###
diff --git a/src/test/modules/test_shmem/test_shmem.c b/src/test/modules/test_shmem/test_shmem.c
index 348f26877af..6adf2405270 100644
--- a/src/test/modules/test_shmem/test_shmem.c
+++ b/src/test/modules/test_shmem/test_shmem.c
@@ -107,6 +107,13 @@ test_shmem_failure_request(void *arg)
 							   .size = (Size) 1024 * 1024 * 1024,
 							   .ptr = &ptr1);
 			break;
+		case 2:
+			ShmemRequestStruct(.name = "test_shmem partial small area",
+							   .size = 1024, .ptr = &ptr1);
+			ShmemRequestStruct(.name = "test_shmem partial huge area",
+							   .size = (Size) 1024 * 1024 * 1024,
+							   .ptr = &ptr1);
+			break;
 		default:
 			elog(ERROR, "unrecognized test_shmem failure mode: %d", failure_mode);
 	}
-- 
2.34.1

