Fix backend state after a failed after-startup shmem request

RegisterShmemCallbacks() left the backend in a bad state, if an error
occurred in the callbacks or if an allocation failed.  Firstly,
'shmem_request_state' was left in wrong state, causing a subsequent
call to RegisterShmemCallbacks() to wrongly take the postmaster
startup codepath or assertion failures in some other functions.
Secondly, the 'pending_shmem_requests' list was not properly cleaned
up, causing a subsequent RegisterShmemCallbacks() to try to process
the stale, already-freed requests.

To fix, add a PG_TRY() block to clean those things up on error.

Author: Ayush Tiwari <ayushtiwari.slg01@gmail.com>
Reviewed-by: Ashutosh Bapat <ashutosh.bapat.oss@gmail.com>
Discussion: https://www.postgresql.org/message-id/CAJTYsWVRRWH48=PcuAo_2Y4Ap6M0QRmzxgUfFkNRtdWK74LjBQ@mail.gmail.com
Backpatch-through: 19

src/backend/storage/ipc/shmem.c

@@ -158,7 +158,9 @@ static List *registered_shmem_callbacks;
 
 /*
  * In the shmem request phase, all the shmem areas requested with the
- * ShmemRequest*() functions are accumulated here.
+ * ShmemRequest*() functions are accumulated in the 'pending_shmem_requests'
+ * list.  The List, the ShmemRequest structs, and the 'options' are all
+ * allocated in TopMemoryContext.
  */
 typedef struct
 {
@@ -166,7 +168,7 @@ typedef struct
 	ShmemRequestKind kind;
 } ShmemRequest;
 
-static List *pending_shmem_requests;
+static List *pending_shmem_requests;	/* List of ShmemRequests */
 
 /*
  * Per-process state machine, for sanity checking that we do things in the
@@ -274,6 +276,7 @@ typedef struct
 static bool firstNumaTouch = true;
 
 static void CallShmemCallbacksAfterStartup(const ShmemCallbacks *callbacks);
+static void ProcessShmemRequestsAfterStartup(const ShmemCallbacks *callbacks);
 static void InitShmemIndexEntry(ShmemRequest *request);
 static bool AttachShmemIndexEntry(ShmemRequest *request, bool missing_ok);
 
@@ -335,6 +338,7 @@ ShmemRequestStructWithOpts(const ShmemStructOpts *options)
 void
 ShmemRequestInternal(ShmemStructOpts *options, ShmemRequestKind kind)
 {
+	MemoryContext oldcontext;
 	ShmemRequest *request;
 
 	/* Check the options */
@@ -374,10 +378,12 @@ ShmemRequestInternal(ShmemStructOpts *options, ShmemRequestKind kind)
 	}
 
 	/* Request looks valid, remember it */
-	request = palloc(sizeof(ShmemRequest));
+	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
+	request = palloc_object(ShmemRequest);
 	request->options = options;
 	request->kind = kind;
 	pending_shmem_requests = lappend(pending_shmem_requests, request);
+	MemoryContextSwitchTo(oldcontext);
 }
 
 /*
@@ -902,27 +908,50 @@ RegisterShmemCallbacks(const ShmemCallbacks *callbacks)
  */
 static void
 CallShmemCallbacksAfterStartup(const ShmemCallbacks *callbacks)
+{
+	Assert(shmem_request_state == SRS_DONE);
+	Assert(pending_shmem_requests == NIL);
+
+	PG_TRY();
+	{
+		shmem_request_state = SRS_REQUESTING;
+
+		/*
+		 * Call the request callback first.  The callback makes
+		 * ShmemRequest*() calls for each shmem area, adding them to
+		 * pending_shmem_requests.
+		 */
+		if (callbacks->request_fn)
+			callbacks->request_fn(callbacks->opaque_arg);
+
+		/* Process all the requests */
+		shmem_request_state = SRS_AFTER_STARTUP_ATTACH_OR_INIT;
+		if (pending_shmem_requests != NIL)
+			ProcessShmemRequestsAfterStartup(callbacks);
+	}
+	PG_FINALLY();
+	{
+		foreach_ptr(ShmemRequest, request, pending_shmem_requests)
+			pfree(request->options);
+		list_free_deep(pending_shmem_requests);
+		pending_shmem_requests = NIL;
+
+		shmem_request_state = SRS_DONE;
+	}
+	PG_END_TRY();
+}
+
+static void
+ProcessShmemRequestsAfterStartup(const ShmemCallbacks *callbacks)
 {
 	bool		found_any;
 	bool		notfound_any;
 
-	Assert(shmem_request_state == SRS_DONE);
-	shmem_request_state = SRS_REQUESTING;
+	/* There should be some requests to process */
+	Assert(pending_shmem_requests != NIL);
 
-	/*
-	 * Call the request callback first.  The callback makes ShmemRequest*()
-	 * calls for each shmem area, adding them to pending_shmem_requests.
-	 */
-	Assert(pending_shmem_requests == NIL);
-	if (callbacks->request_fn)
-		callbacks->request_fn(callbacks->opaque_arg);
-	shmem_request_state = SRS_AFTER_STARTUP_ATTACH_OR_INIT;
-
-	if (pending_shmem_requests == NIL)
-	{
-		shmem_request_state = SRS_DONE;
-		return;
-	}
+	/* Caller manages the global state variable */
+	Assert(shmem_request_state == SRS_AFTER_STARTUP_ATTACH_OR_INIT);
 
 	/*
 	 * Hold ShmemIndexLock while we allocate all the shmem entries and run all
@@ -940,7 +969,11 @@ CallShmemCallbacksAfterStartup(const ShmemCallbacks *callbacks)
 	found_any = notfound_any = false;
 	foreach_ptr(ShmemRequest, request, pending_shmem_requests)
 	{
-		if (hash_search(ShmemIndex, request->options->name, HASH_FIND, NULL))
+		ShmemIndexEnt *index_entry;
+
+		index_entry = (ShmemIndexEnt *)
+			hash_search(ShmemIndex, request->options->name, HASH_FIND, NULL);
+		if (index_entry)
 			found_any = true;
 		else
 			notfound_any = true;
@@ -958,11 +991,7 @@ CallShmemCallbacksAfterStartup(const ShmemCallbacks *callbacks)
 			AttachShmemIndexEntry(request, false);
 		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)
@@ -977,7 +1006,6 @@ CallShmemCallbacksAfterStartup(const ShmemCallbacks *callbacks)
 	}
 
 	LWLockRelease(ShmemIndexLock);
-	shmem_request_state = SRS_DONE;
 }
 
 /*

src/test/modules/test_shmem/t/001_late_shmem_alloc.pl

@@ -78,5 +78,33 @@ else
 	);
 }
 
+# clean up
 $node->stop;
+$node->adjust_conf('postgresql.conf', "shared_preload_libraries", undef);
+
+###
+# Test "out of shared memory" in an after-startup request
+###
+$node->start;
+my $session = $node->background_psql('postgres', on_error_stop => 0);
+
+# make the request larger than the memory reserved for after-startup
+# requests.
+$session->query(q[SET test_shmem.area_size = '128kB';]);
+
+$session->query("SELECT get_test_shmem_attach_count();");
+like(
+	$session->{stderr},
+	qr/not enough shared memory/,
+	"an after-startup request larger than the reserve fails");
+
+# The server and the backend keep running.  Since only one area was
+# requested, it gets cleaned up on allocation failure.  Verify that a
+# request for a smaller area succeeds in the same session.
+$session->{stderr} = '';
+$session->query("SET test_shmem.area_size = default;");
+$session->query_safe("SELECT get_test_shmem_attach_count();");
+$session->quit;
+$node->stop;
+
 done_testing();

src/test/modules/test_shmem/test_shmem.c

@@ -20,20 +20,28 @@
 #include "fmgr.h"
 #include "miscadmin.h"
 #include "storage/shmem.h"
+#include "utils/guc.h"
+#include "utils/injection_point.h"
 
 
 PG_MODULE_MAGIC;
 
 typedef struct TestShmemData
 {
-	int			value;
 	bool		initialized;
 	int			attach_count;
+	char		dummy_data[FLEXIBLE_ARRAY_MEMBER];
 } TestShmemData;
 
 static TestShmemData *TestShmem;
 
+#define MIN_TEST_AREA_BYTES sizeof(TestShmemData)
+#define DEFAULT_TEST_AREA_BYTES MIN_TEST_AREA_BYTES
+#define MAX_TEST_AREA_BYTES 1000000
+
 static bool attached_or_initialized = false;
+static int	test_shmem_area_size = MIN_TEST_AREA_BYTES;
+static bool test_shmem_guc_defined = false;
 
 static void test_shmem_request(void *arg);
 static void test_shmem_init(void *arg);
@@ -52,7 +60,7 @@ test_shmem_request(void *arg)
 	elog(LOG, "test_shmem_request callback called");
 
 	ShmemRequestStruct(.name = "test_shmem area",
-					   .size = sizeof(TestShmemData),
+					   .size = test_shmem_area_size,
 					   .ptr = (void **) &TestShmem);
 }
 
@@ -86,6 +94,22 @@ void
 _PG_init(void)
 {
 	elog(LOG, "test_shmem module's _PG_init called");
+
+	if (!test_shmem_guc_defined)
+	{
+		DefineCustomIntVariable("test_shmem.area_size",
+								"Size of the shmem area to request.",
+								NULL,
+								&test_shmem_area_size,
+								DEFAULT_TEST_AREA_BYTES,
+								MIN_TEST_AREA_BYTES,
+								MAX_TEST_AREA_BYTES,
+								PGC_USERSET,
+								GUC_UNIT_BYTE,
+								NULL, NULL, NULL);
+		MarkGUCPrefixReserved("test_shmem");
+		test_shmem_guc_defined = true;
+	}
 	RegisterShmemCallbacks(&TestShmemCallbacks);
 }