From d06792297a82fb0a0a1d8dbb1f4064c324cb485a Mon Sep 17 00:00:00 2001
From: Yuriy Grigoryev <ju.grigorev@ftdata.ru>
Date: Tue, 8 Sep 2026 13:30:18 +0700
Subject: [PATCH] Allocate pgstats entry body before shared hash insert

8191e0c made pgstat_init_entry() return NULL on DSA allocation failure
so callers could delete a just-inserted hash entry.  That works when
dsa_allocate_extended(..., DSA_ALLOC_NO_OOM) returns InvalidDsaPointer,
but not when creating a new DSM segment raises ERROR (for example ENOSPC
on posix shared memory).

pgstat_init_entry() marked the hash entry live before allocating the
body.  If dsm_create() failed after the insert, error unwind released
the dshash lock but left a visible entry with body == InvalidDsaPointer.
The next backend crashed in pgstat_acquire_entry_ref().

Split allocation from hash initialization: allocate the DSA chunk first,
insert the hash entry only with a valid body, and free the chunk if a
concurrent insert wins or the insert returns NULL.  Both callers use
dshash_find_or_insert_extended() with DSHASH_INSERT_NO_OOM so that path
can free the preallocated body.

Discussion: https://postgr.es/m/ddc3ecfb01ce4e9698b23cc59767f016@localhost.localdomain
---
 src/backend/utils/activity/pgstat.c       |  45 ++++++----
 src/backend/utils/activity/pgstat_shmem.c | 105 ++++++++++++----------
 src/include/utils/pgstat_internal.h       |   4 +-
 3 files changed, 91 insertions(+), 63 deletions(-)

diff --git a/src/backend/utils/activity/pgstat.c b/src/backend/utils/activity/pgstat.c
index 5177f88..3f4fe60 100644
--- a/src/backend/utils/activity/pgstat.c
+++ b/src/backend/utils/activity/pgstat.c
@@ -2007,6 +2007,7 @@ pgstat_read_statsfile(void)
 					PgStatShared_HashEntry *p;
 					PgStatShared_Common *header;
 					const PgStat_KindInfo *kind_info = NULL;
+					dsa_pointer chunk;
 
 					CHECK_FOR_INTERRUPTS();
 
@@ -2095,12 +2096,39 @@ pgstat_read_statsfile(void)
 					 * This intentionally doesn't use pgstat_get_entry_ref() -
 					 * putting all stats into checkpointer's
 					 * pgStatEntryRefHash would be wasted effort and memory.
+					 *
+					 * Allocate the DSA body before inserting the hash entry,
+					 * so a dsm_create failure cannot leave a half-initialized
+					 * shared entry behind.
 					 */
-					p = dshash_find_or_insert(pgStatLocal.shared_hash, &key, &found);
+					chunk = pgstat_alloc_entry_body(key.kind);
+					if (chunk == InvalidDsaPointer)
+					{
+						/*
+						 * It would be tempting to switch this ERROR to a
+						 * WARNING, but it would mean that all the statistics
+						 * are discarded when the environment fails on OOM.
+						 */
+						elog(ERROR, "could not allocate entry %u/%u/%" PRIu64 " of type %c",
+							 key.kind, key.dboid,
+							 key.objid, t);
+					}
+
+					p = dshash_find_or_insert_extended(pgStatLocal.shared_hash,
+													   &key, &found,
+													   DSHASH_INSERT_NO_OOM);
+					if (!p)
+					{
+						dsa_free(pgStatLocal.dsa, chunk);
+						elog(ERROR, "could not insert entry %u/%u/%" PRIu64 " of type %c",
+							 key.kind, key.dboid,
+							 key.objid, t);
+					}
 
 					/* don't allow duplicate entries */
 					if (found)
 					{
+						dsa_free(pgStatLocal.dsa, chunk);
 						dshash_release_lock(pgStatLocal.shared_hash, p);
 						elog(WARNING, "found duplicate stats entry %u/%u/%" PRIu64 " of type %c",
 							 key.kind, key.dboid,
@@ -2108,20 +2136,7 @@ pgstat_read_statsfile(void)
 						goto error;
 					}
 
-					header = pgstat_init_entry(key.kind, p);
-					if (header == NULL)
-					{
-						dshash_delete_entry(pgStatLocal.shared_hash, p);
-
-						/*
-						 * It would be tempting to switch this ERROR to a
-						 * WARNING, but it would mean that all the statistics
-						 * are discarded when the environment fails on OOM.
-						 */
-						elog(ERROR, "could not allocate entry %u/%u/%" PRIu64 " of type %c",
-							 key.kind, key.dboid,
-							 key.objid, t);
-					}
+					header = pgstat_init_entry(key.kind, p, chunk);
 					dshash_release_lock(pgStatLocal.shared_hash, p);
 
 					if (!read_chunk(fpin,
diff --git a/src/backend/utils/activity/pgstat_shmem.c b/src/backend/utils/activity/pgstat_shmem.c
index d8ac9d6..bf886b4 100644
--- a/src/backend/utils/activity/pgstat_shmem.c
+++ b/src/backend/utils/activity/pgstat_shmem.c
@@ -303,53 +303,58 @@ pgstat_detach_shmem(void)
  */
 
 /*
- * Initialize entry newly-created.
+ * Allocate the DSA body for a new variable-numbered pgstats entry.
  *
- * Returns NULL in the event of an allocation failure, so as callers can
- * take cleanup actions as the entry initialized is already inserted in the
- * shared hashtable.
+ * Returns InvalidDsaPointer if the allocation fails without throwing.  Call
+ * this before inserting a hash entry: dsa_allocate_extended() can still raise
+ * ERROR when creating a new DSM segment (for example ENOSPC), and doing that
+ * after the insert would leave a live hash entry with body ==
+ * InvalidDsaPointer.
+ */
+dsa_pointer
+pgstat_alloc_entry_body(PgStat_Kind kind)
+{
+	const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
+
+	return dsa_allocate_extended(pgStatLocal.dsa,
+								 kind_info->shared_size,
+								 DSA_ALLOC_ZERO | DSA_ALLOC_NO_OOM);
+}
+
+/*
+ * Initialize a newly-inserted hash entry around an already-allocated DSA
+ * body.
+ *
+ * The caller must hold the dshash partition lock.  The entry cannot be found
+ * by other backends until that lock is released, so it is safe to publish
+ * refcount/dropped/body here.  Caller needs to increment the refcount further
+ * if a longer-lived reference is needed.
+ *
+ * chunk must be a valid pointer from pgstat_alloc_entry_body().
  */
 PgStatShared_Common *
 pgstat_init_entry(PgStat_Kind kind,
-				  PgStatShared_HashEntry *shhashent)
+				  PgStatShared_HashEntry *shhashent,
+				  dsa_pointer chunk)
 {
-	/* Create new stats entry. */
-	dsa_pointer chunk;
 	PgStatShared_Common *shheader;
 	const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
 
-	/*
-	 * Initialize refcount to 1, marking it as valid / not dropped. The entry
-	 * can't be freed before the initialization because it can't be found as
-	 * long as we hold the dshash partition lock. Caller needs to increase
-	 * further if a longer lived reference is needed.
-	 */
-	pg_atomic_init_u32(&shhashent->refcount, 1);
-
-	/*
-	 * Initialize "generation" to 0, as freshly created.
-	 */
-	pg_atomic_init_u32(&shhashent->generation, 0);
-	shhashent->dropped = false;
-
-	chunk = dsa_allocate_extended(pgStatLocal.dsa,
-								  kind_info->shared_size,
-								  DSA_ALLOC_ZERO | DSA_ALLOC_NO_OOM);
-	if (chunk == InvalidDsaPointer)
-		return NULL;
+	Assert(DsaPointerIsValid(chunk));
 
 	shheader = dsa_get_address(pgStatLocal.dsa, chunk);
 	shheader->magic = 0xdeadbeef;
+	LWLockInitialize(&shheader->lock, LWTRANCHE_PGSTATS_DATA);
 
-	/* Link the new entry from the hash entry. */
+	pg_atomic_init_u32(&shhashent->refcount, 1);
+	pg_atomic_init_u32(&shhashent->generation, 0);
+	shhashent->dropped = false;
 	shhashent->body = chunk;
 
 	/* Increment entry count, if required. */
 	if (kind_info->track_entry_count)
 		pg_atomic_fetch_add_u64(&pgStatLocal.shmem->entry_counts[kind - 1], 1);
 
-	LWLockInitialize(&shheader->lock, LWTRANCHE_PGSTATS_DATA);
-
 	return shheader;
 }
 
@@ -545,6 +550,24 @@ pgstat_get_entry_ref(PgStat_Kind kind, Oid dboid, uint64 objid, bool create,
 	if (create && !shhashent)
 	{
 		bool		shfound;
+		dsa_pointer chunk;
+
+		/*
+		 * Allocate the stats body before inserting a hash entry.  Creating a
+		 * new DSA segment can raise ERROR (e.g. ENOSPC on posix shm); doing
+		 * that after the insert would leave a live hash entry with an
+		 * invalid body.
+		 */
+		chunk = pgstat_alloc_entry_body(kind);
+		if (chunk == InvalidDsaPointer)
+		{
+			pgstat_release_entry_ref(key, entry_ref, false);
+			ereport(ERROR,
+					(errcode(ERRCODE_OUT_OF_MEMORY),
+					 errmsg("out of memory"),
+					 errdetail("Failed while allocating entry %u/%u/%" PRIu64 ".",
+							   key.kind, key.dboid, key.objid)));
+		}
 
 		/*
 		 * It's possible that somebody created the entry since the above
@@ -556,6 +579,8 @@ pgstat_get_entry_ref(PgStat_Kind kind, Oid dboid, uint64 objid, bool create,
 												   DSHASH_INSERT_NO_OOM);
 		if (!shhashent)
 		{
+			dsa_free(pgStatLocal.dsa, chunk);
+
 			/*
 			 * Clean up the local reference when failing insert into the
 			 * shared hashtable.
@@ -570,24 +595,7 @@ pgstat_get_entry_ref(PgStat_Kind kind, Oid dboid, uint64 objid, bool create,
 
 		if (!shfound)
 		{
-			shheader = pgstat_init_entry(kind, shhashent);
-			if (shheader == NULL)
-			{
-				/*
-				 * 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,
-						(errcode(ERRCODE_OUT_OF_MEMORY),
-						 errmsg("out of memory"),
-						 errdetail("Failed while allocating entry %u/%u/%" PRIu64 ".",
-								   key.kind, key.dboid, key.objid)));
-			}
+			shheader = pgstat_init_entry(kind, shhashent, chunk);
 			pgstat_acquire_entry_ref(entry_ref, shhashent, shheader);
 
 			if (created_entry != NULL)
@@ -595,6 +603,9 @@ pgstat_get_entry_ref(PgStat_Kind kind, Oid dboid, uint64 objid, bool create,
 
 			return entry_ref;
 		}
+
+		/* Concurrent insert won; drop the unused body. */
+		dsa_free(pgStatLocal.dsa, chunk);
 	}
 
 	if (!shhashent)
diff --git a/src/include/utils/pgstat_internal.h b/src/include/utils/pgstat_internal.h
index 14369e5..201e572 100644
--- a/src/include/utils/pgstat_internal.h
+++ b/src/include/utils/pgstat_internal.h
@@ -839,8 +839,10 @@ extern void pgstat_reset_matching_entries(bool (*do_reset) (PgStatShared_HashEnt
 										  TimestampTz ts);
 
 extern void pgstat_request_entry_refs_gc(void);
+extern dsa_pointer pgstat_alloc_entry_body(PgStat_Kind kind);
 extern PgStatShared_Common *pgstat_init_entry(PgStat_Kind kind,
-											  PgStatShared_HashEntry *shhashent);
+											  PgStatShared_HashEntry *shhashent,
+											  dsa_pointer chunk);
 
 
 /*
-- 
2.50.1 (Apple Git-155)

