Subject: [RFC PATCH] Clean up failed DSM attachments and incomplete pgstats refs

This is a discussion patch against REL_18_STABLE at
13a9be148e530b9b5c0af2b499e3a08c1f5eb053.

An ERROR while mapping a DSM segment can occur after dsm_attach() has
created a backend-local descriptor and incremented the shared reference
count.  If the caller deliberately uses no ResourceOwner, as a pinned DSA
does, that descriptor survives when the ERROR is caught.  A retry then fails
with "can't attach the same segment more than once".

An ERROR after pgstat_get_entry_ref_cached() has inserted a local cache entry
can similarly leave an incomplete PgStat_EntryRef.  A later cache lookup can
repair it, but pgstat_gc_entry_refs() currently runs before that lookup and
unconditionally dereferences shared_entry.

Make dsm_attach() exception-safe, finish PgStat_EntryRef initialization
before releasing the dshash lock, remove the local reference on the explicit
DSA allocation failure path, and let pgstats GC discard incomplete local
entries.  Also include the substance of master commit 16ebc196 for the
adjacent local-allocation failure; that hunk should be omitted on master.

No automated fault-injection test is included in this RFC.

diff --git a/src/backend/storage/ipc/dsm.c b/src/backend/storage/ipc/dsm.c
--- a/src/backend/storage/ipc/dsm.c
+++ b/src/backend/storage/ipc/dsm.c
@@ -742,8 +742,25 @@ dsm_attach(dsm_handle h)
 
 	/* Here's where we actually try to map the segment. */
 	if (!is_main_region_dsm_handle(seg->handle))
-		dsm_impl_op(DSM_OP_ATTACH, seg->handle, 0, &seg->impl_private,
-					&seg->mapped_address, &seg->mapped_size, ERROR);
+	{
+		/*
+		 * dsm_impl_op() can raise an error after we have created the
+		 * backend-local descriptor and incremented the shared reference
+		 * count.  A ResourceOwner normally cleans that up during abort, but
+		 * session-pinned mappings have no ResourceOwner.
+		 */
+		PG_TRY();
+		{
+			dsm_impl_op(DSM_OP_ATTACH, seg->handle, 0, &seg->impl_private,
+						&seg->mapped_address, &seg->mapped_size, ERROR);
+		}
+		PG_CATCH();
+		{
+			dsm_detach(seg);
+			PG_RE_THROW();
+		}
+		PG_END_TRY();
+	}
 
 	return seg;
 }
@@ -836,6 +853,6 @@ dsm_detach(dsm_segment *seg)
 	 * works, because retrying is likely to fail in the same way.
 	 */
-	if (seg->mapped_address != NULL)
+	if (seg->mapped_address != NULL || seg->impl_private != NULL)
 	{
 		if (!is_main_region_dsm_handle(seg->handle))
 			dsm_impl_op(DSM_OP_DETACH, seg->handle, 0, &seg->impl_private,
diff --git a/src/backend/utils/activity/pgstat_shmem.c b/src/backend/utils/activity/pgstat_shmem.c
--- a/src/backend/utils/activity/pgstat_shmem.c
+++ b/src/backend/utils/activity/pgstat_shmem.c
@@ -387,11 +387,16 @@ pgstat_acquire_entry_ref(PgStat_EntryRef *entry_ref,
 
 	pg_atomic_fetch_add_u32(&shhashent->refcount, 1);
 
-	dshash_release_lock(pgStatLocal.shared_hash, shhashent);
-
 	entry_ref->shared_stats = shheader;
 	entry_ref->shared_entry = shhashent;
 	entry_ref->generation = pg_atomic_read_u32(&shhashent->generation);
+
+	/*
+	 * Complete the local reference before releasing the lock.  Releasing an
+	 * LWLock can process a pending interrupt, and callers may catch the
+	 * resulting error and continue using the backend-local cache.
+	 */
+	dshash_release_lock(pgStatLocal.shared_hash, shhashent);
 }
 
 /*
@@ -415,9 +420,23 @@ pgstat_get_entry_ref_cached(PgStat_HashKey key, PgStat_EntryRef **entry_ref_p)
 	{
 		PgStat_EntryRef *entry_ref;
 
-		cache_entry->entry_ref = entry_ref =
-			MemoryContextAlloc(pgStatSharedRefContext,
-							   sizeof(PgStat_EntryRef));
+		entry_ref = MemoryContextAllocExtended(pgStatSharedRefContext,
+										   sizeof(PgStat_EntryRef),
+										   MCXT_ALLOC_NO_OOM);
+		if (unlikely(entry_ref == NULL))
+		{
+			/*
+			 * Clean the hash entry to keep the table consistent in the
+			 * backend.
+			 */
+			pgstat_entry_ref_hash_delete(pgStatEntryRefHash, key);
+
+			ereport(ERROR,
+					(errcode(ERRCODE_OUT_OF_MEMORY),
+					 errmsg("out of memory")));
+		}
+
+		cache_entry->entry_ref = entry_ref;
 		entry_ref->shared_stats = NULL;
 		entry_ref->shared_entry = NULL;
 		entry_ref->pending = NULL;
@@ -525,9 +544,12 @@ pgstat_get_entry_ref(PgStat_Kind kind, Oid dboid, uint64 objid, bool create,
 			if (shheader == NULL)
 			{
 				/*
-				 * Failed the allocation of a new entry, so clean up the
-				 * shared hashtable before giving up.
+				 * Failed the allocation of a new entry, so clean up both
+				 * the local reference and the shared hashtable before giving
+				 * up.  Clean the local state first, since releasing the
+				 * dshash lock can process a pending interrupt.
 				 */
+				pgstat_release_entry_ref(key, entry_ref, false);
 				dshash_delete_entry(pgStatLocal.shared_hash, shhashent);
 
 				ereport(ERROR,
@@ -773,6 +795,20 @@ pgstat_gc_entry_refs(void)
 	{
 		PgStat_EntryRef *entry_ref = ent->entry_ref;
 
+		/*
+		 * An error during a cache miss can leave behind a local hash entry
+		 * that has not yet acquired a shared reference.  Such an entry is
+		 * safe to discard, but must not be dereferenced below.
+		 */
+		if (entry_ref == NULL || entry_ref->shared_stats == NULL)
+		{
+			Assert(entry_ref == NULL ||
+				   (entry_ref->shared_entry == NULL &&
+					entry_ref->pending == NULL));
+			pgstat_release_entry_ref(ent->key, entry_ref, false);
+			continue;
+		}
+
 		Assert(!entry_ref->shared_stats ||
 			   entry_ref->shared_stats->magic == 0xdeadbeef);