Track which shmem areas have been fully initialized

If SHMEM_CALLBACKS_ALLOW_AFTER_STARTUP is used to allocate shared
memory after startup, but the initialization fails half-way through,
the shmem area is left in an indeterminate state.  Furthermore, if
multiple shmem areas are registered in one RegisterShmemCallbacks()
call, some might be allocated while others are not.

This commit adds an explicit 'initialized' flag to each shmem area.
We still leave behind an uninitialized area on error, but at least
they are now clearly marked, and you get a slightly nicer error
message if you try to re-register them.  It'd be nice to clean up more
thoroughly and support actually retrying the allocations, but in
practice, the most likely reason for a shmem allocation or
initialization to fail is that you are out of shared memory and
retrying wouldn't help with that.

This isn't exactly a new problem, the old ShmemInitStruct() interface
had similar issues if the initialization code failed, or if you
allocated multiple structs and some allocations failed.  It was just
left to the calling code to deal with it.

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

doc/src/sgml/xfunc.sgml

@@ -3740,7 +3740,10 @@ my_shmem_init(void *arg)
       on whether the requested memory areas were already initialized by
       another backend. The callbacks will be called while holding an internal
       lock (ShmemIndexLock), which prevents the race condition of two backends
-      trying to initialize the memory area at the same time.
+      trying to initialize the memory area at the same time.  If the
+      allocation or initialization fails for any reason, the shared memory
+      areas are left in an abandoned state and any attempt to attach or
+      re-initialize them will fail until the server is restarted.
      </para>
     </sect3>
 

src/backend/storage/ipc/shmem.c

@@ -144,6 +144,8 @@
 #include "utils/builtins.h"
 #include "utils/tuplestore.h"
 
+typedef struct ShmemIndexEnt ShmemIndexEnt;
+
 /*
  * Registered callbacks.
  *
@@ -166,6 +168,9 @@ typedef struct
 {
 	ShmemStructOpts *options;
 	ShmemRequestKind kind;
+
+	/* InitShmemIndexEntry() sets this pointer when the area is allocated */
+	ShmemIndexEnt *index_entry;
 } ShmemRequest;
 
 static List *pending_shmem_requests;	/* List of ShmemRequests */
@@ -264,12 +269,13 @@ static HTAB *ShmemIndex;
 #define SHMEM_INDEX_ADDITIONAL_SIZE		 (128)
 
 /* this is a hash bucket in the shmem index table */
-typedef struct
+typedef struct ShmemIndexEnt
 {
 	char		key[SHMEM_INDEX_KEYSIZE];	/* string name */
 	void	   *location;		/* location in shared mem */
 	Size		size;			/* # bytes requested for the structure */
 	Size		allocated_size; /* # bytes actually allocated */
+	bool		initialized;	/* has the init callback been run? */
 } ShmemIndexEnt;
 
 /* To get reliable results for NUMA inquiry we need to "touch pages" once */
@@ -382,6 +388,7 @@ ShmemRequestInternal(ShmemStructOpts *options, ShmemRequestKind kind)
 	request = palloc_object(ShmemRequest);
 	request->options = options;
 	request->kind = kind;
+	request->index_entry = NULL;
 	pending_shmem_requests = lappend(pending_shmem_requests, request);
 	MemoryContextSwitchTo(oldcontext);
 }
@@ -441,10 +448,7 @@ ShmemInitRequested(void)
 	foreach_ptr(ShmemRequest, request, pending_shmem_requests)
 	{
 		InitShmemIndexEntry(request);
-		pfree(request->options);
 	}
-	list_free_deep(pending_shmem_requests);
-	pending_shmem_requests = NIL;
 
 	/*
 	 * Call the subsystem-specific init callbacks to finish initialization of
@@ -456,6 +460,15 @@ ShmemInitRequested(void)
 			callbacks->init_fn(callbacks->opaque_arg);
 	}
 
+	/* Now we can mark all the areas as initialized and free the requests */
+	foreach_ptr(ShmemRequest, request, pending_shmem_requests)
+	{
+		request->index_entry->initialized = true;
+		pfree(request->options);
+	}
+	list_free_deep(pending_shmem_requests);
+	pending_shmem_requests = NIL;
+
 	shmem_request_state = SRS_DONE;
 }
 
@@ -557,7 +570,12 @@ InitShmemIndexEntry(ShmemRequest *request)
 	index_entry->allocated_size = allocated_size;
 	index_entry->location = structPtr;
 
-	/* Initialize depending on the kind of shmem area it is */
+	/*
+	 * The area is considered fully initialized only after the subsystem's
+	 * init callback has been called.  For now, perform only basic
+	 * initialization based on the kind of shmem area it is.
+	 */
+	index_entry->initialized = false;
 	switch (request->kind)
 	{
 		case SHMEM_KIND_STRUCT:
@@ -571,6 +589,9 @@ InitShmemIndexEntry(ShmemRequest *request)
 			shmem_slru_init(structPtr, request->options);
 			break;
 	}
+
+	/* return the pointer to the entry to the caller */
+	request->index_entry = index_entry;
 }
 
 /*
@@ -600,6 +621,20 @@ AttachShmemIndexEntry(ShmemRequest *request, bool missing_ok)
 		return false;
 	}
 
+	/*
+	 * If it was previously allocated but not fully initialized, error out.
+	 * There is currently no way of retrying or cleaning up an uninitialized
+	 * entry, it just lingers until the server is shut down.  But this can
+	 * only happen when allocating areas after postmaster startup, and it's
+	 * unlikely that you could successfully retry anyway.  The most likely
+	 * reason for failed initialization is that you are out of shared memory
+	 * and retrying won't help with that.
+	 */
+	if (!index_entry->initialized)
+		ereport(ERROR,
+				(errmsg("cannot attach to shared memory struct \"%s\" because it was not fully initialized",
+						request->options->name)));
+
 	/* Check that the size in the index matches the request */
 	if (index_entry->size != request->options->size &&
 		request->options->size != SHMEM_ATTACH_UNKNOWN_SIZE)
@@ -628,6 +663,8 @@ AttachShmemIndexEntry(ShmemRequest *request, bool missing_ok)
 			break;
 	}
 
+	request->index_entry = index_entry;
+
 	return true;
 }
 
@@ -738,6 +775,7 @@ InitShmemAllocator(PGShmemHeader *seghdr)
 		result->size = ShmemAllocator->index_size;
 		result->allocated_size = ShmemAllocator->index_size;
 		result->location = ShmemAllocator->index;
+		result->initialized = true;
 	}
 }
 
@@ -974,7 +1012,17 @@ ProcessShmemRequestsAfterStartup(const ShmemCallbacks *callbacks)
 		index_entry = (ShmemIndexEnt *)
 			hash_search(ShmemIndex, request->options->name, HASH_FIND, NULL);
 		if (index_entry)
+		{
+			/*
+			 * Check for a half-initialized area.  (See also similar check in
+			 * AttachShmemIndexEntry())
+			 */
+			if (!index_entry->initialized)
+				ereport(ERROR,
+						(errmsg("cannot attach to shared memory struct \"%s\" because it was not fully initialized",
+								request->options->name)));
 			found_any = true;
+		}
 		else
 			notfound_any = true;
 	}
@@ -1005,6 +1053,11 @@ ProcessShmemRequestsAfterStartup(const ShmemCallbacks *callbacks)
 			callbacks->init_fn(callbacks->opaque_arg);
 	}
 
+	foreach_ptr(ShmemRequest, request, pending_shmem_requests)
+	{
+		request->index_entry->initialized = true;
+	}
+
 	LWLockRelease(ShmemIndexLock);
 }
 
@@ -1069,7 +1122,11 @@ ShmemInitStruct(const char *name, Size size, bool *foundPtr)
 
 	/* Initialize it if not found */
 	if (!*foundPtr)
+	{
 		InitShmemIndexEntry(&request);
+		/* no additional initialization needed */
+		request.index_entry->initialized = true;
+	}
 
 	LWLockRelease(ShmemIndexLock);
 

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) \

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',
     ],

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

@@ -7,17 +7,20 @@ use PostgreSQL::Test::Cluster;
 use PostgreSQL::Test::Utils;
 use Test::More;
 
+# Initialize a cluster with the extension installed.  The tests will
+# call the function that comes with the extension to load it.
+my $node = PostgreSQL::Test::Cluster->new('main');
+$node->init;
+$node->start;
+$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
 ###
-my $node = PostgreSQL::Test::Cluster->new('main');
-$node->init;
 $node->start;
 
-
-$node->safe_psql("postgres", "CREATE EXTENSION test_shmem;");
-
 # Check that the attach counter is incremented on a new connection
 my $attach_count1 =
   $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();");
@@ -25,6 +28,7 @@ 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;
 
 ###
@@ -82,6 +86,42 @@ else
 $node->stop;
 $node->adjust_conf('postgresql.conf', "shared_preload_libraries", undef);
 
+###
+# Test a failure in initializing the shared memory area
+###
+SKIP:
+{
+	skip "injection points not supported by this build",
+	  if $ENV{enable_injection_points} ne 'yes';
+	$node->start;
+	$node->safe_psql("postgres", "CREATE EXTENSION injection_points;");
+	$node->safe_psql("postgres",
+		"SELECT injection_points_attach('test-shmem-init', 'error');");
+
+	# Try to load the extension library. It will hit the injected
+	# error in the init callback.
+	my (undef, undef, $stderr) =
+	  $node->psql("postgres", "SELECT get_test_shmem_attach_count();");
+	like(
+		$stderr,
+		qr/error triggered for injection point test-shmem-init/,
+		"failure in initialization is reported");
+	$node->safe_psql("postgres",
+		"SELECT injection_points_detach('test-shmem-init');");
+
+	# The error leaves the shared memory area in a broken state.
+	# Attempting to initialize or attach it again will fail, until the
+	# server is restarted.
+	(undef, undef, $stderr) =
+	  $node->psql("postgres", "SELECT get_test_shmem_attach_count();");
+	like(
+		$stderr,
+		qr/cannot attach to shared memory/,
+		"post-init extension creation fails");
+
+	$node->stop;
+}
+
 ###
 # Test "out of shared memory" in an after-startup request
 ###

src/test/modules/test_shmem/test_shmem.c

@@ -68,6 +68,9 @@ static void
 test_shmem_init(void *arg)
 {
 	elog(LOG, "init callback called");
+
+	INJECTION_POINT("test-shmem-init", NULL);
+
 	if (TestShmem->initialized)
 		elog(ERROR, "shmem area already initialized");
 	TestShmem->initialized = true;