From 295352adb54faf36fd730922ada032a4bbaaa941 Mon Sep 17 00:00:00 2001
From: Mikhail Kot <to@myrrc.dev>
Date: Tue, 2 Sep 2025 21:02:04 +0100
Subject: [PATCH] fix sshashent access segfault when it's half initialized due
 to OOM

---
 src/backend/utils/activity/pgstat_database.c     |  3 +++
 src/backend/utils/activity/pgstat_function.c     |  3 +++
 src/backend/utils/activity/pgstat_relation.c     |  3 +++
 src/backend/utils/activity/pgstat_shmem.c        | 13 ++++++++++---
 src/backend/utils/activity/pgstat_subscription.c |  3 +++
 5 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/backend/utils/activity/pgstat_database.c b/src/backend/utils/activity/pgstat_database.c
index b31f20d41bc..71e897d684b 100644
--- a/src/backend/utils/activity/pgstat_database.c
+++ b/src/backend/utils/activity/pgstat_database.c
@@ -430,6 +430,9 @@ pgstat_database_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 	PgStatShared_Database *sharedent;
 	PgStat_StatDBEntry *pendingent;
 
+	if (!entry_ref)
+		return false;
+
 	pendingent = (PgStat_StatDBEntry *) entry_ref->pending;
 	sharedent = (PgStatShared_Database *) entry_ref->shared_stats;
 
diff --git a/src/backend/utils/activity/pgstat_function.c b/src/backend/utils/activity/pgstat_function.c
index 6214f93d36e..02916b80cb2 100644
--- a/src/backend/utils/activity/pgstat_function.c
+++ b/src/backend/utils/activity/pgstat_function.c
@@ -195,6 +195,9 @@ pgstat_function_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 	PgStat_FunctionCounts *localent;
 	PgStatShared_Function *shfuncent;
 
+	if (!entry_ref)
+		return false;
+
 	localent = (PgStat_FunctionCounts *) entry_ref->pending;
 	shfuncent = (PgStatShared_Function *) entry_ref->shared_stats;
 
diff --git a/src/backend/utils/activity/pgstat_relation.c b/src/backend/utils/activity/pgstat_relation.c
index 69df741cbf6..20dfa137ba6 100644
--- a/src/backend/utils/activity/pgstat_relation.c
+++ b/src/backend/utils/activity/pgstat_relation.c
@@ -821,6 +821,9 @@ pgstat_relation_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 	PgStat_StatTabEntry *tabentry;	/* table entry of shared stats */
 	PgStat_StatDBEntry *dbentry;	/* pending database entry */
 
+	if (!entry_ref)
+		return false;
+
 	dboid = entry_ref->shared_entry->key.dboid;
 	lstats = (PgStat_TableStatus *) entry_ref->pending;
 	shtabstats = (PgStatShared_Relation *) entry_ref->shared_stats;
diff --git a/src/backend/utils/activity/pgstat_shmem.c b/src/backend/utils/activity/pgstat_shmem.c
index 62de3474453..76b12b0b1b4 100644
--- a/src/backend/utils/activity/pgstat_shmem.c
+++ b/src/backend/utils/activity/pgstat_shmem.c
@@ -311,7 +311,12 @@ pgstat_init_entry(PgStat_Kind kind,
 	pg_atomic_init_u32(&shhashent->generation, 0);
 	shhashent->dropped = false;
 
-	chunk = dsa_allocate0(pgStatLocal.dsa, pgstat_get_kind_info(kind)->shared_size);
+	chunk = dsa_allocate_extended(pgStatLocal.dsa,
+								  pgstat_get_kind_info(kind)->shared_size,
+								  DSA_ALLOC_ZERO | DSA_ALLOC_NO_OOM);
+	if (!chunk)
+		return NULL;
+
 	shheader = dsa_get_address(pgStatLocal.dsa, chunk);
 	shheader->magic = 0xdeadbeef;
 
@@ -508,7 +513,9 @@ pgstat_get_entry_ref(PgStat_Kind kind, Oid dboid, uint64 objid, bool create,
 		shhashent = dshash_find_or_insert(pgStatLocal.shared_hash, &key, &shfound);
 		if (!shfound)
 		{
-			shheader = pgstat_init_entry(kind, shhashent);
+			if ((shheader = pgstat_init_entry(kind, shhashent)) == NULL)
+				return NULL;
+
 			pgstat_acquire_entry_ref(entry_ref, shhashent, shheader);
 
 			if (created_entry != NULL)
@@ -699,7 +706,7 @@ pgstat_get_entry_ref_locked(PgStat_Kind kind, Oid dboid, uint64 objid,
 	entry_ref = pgstat_get_entry_ref(kind, dboid, objid, true, NULL);
 
 	/* lock the shared entry to protect the content, skip if failed */
-	if (!pgstat_lock_entry(entry_ref, nowait))
+	if (!entry_ref || !pgstat_lock_entry(entry_ref, nowait))
 		return NULL;
 
 	return entry_ref;
diff --git a/src/backend/utils/activity/pgstat_subscription.c b/src/backend/utils/activity/pgstat_subscription.c
index f9a1c831a07..209a37607e4 100644
--- a/src/backend/utils/activity/pgstat_subscription.c
+++ b/src/backend/utils/activity/pgstat_subscription.c
@@ -105,6 +105,9 @@ pgstat_subscription_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 	PgStat_BackendSubEntry *localent;
 	PgStatShared_Subscription *shsubent;
 
+	if (!entry_ref)
+		return false;
+
 	localent = (PgStat_BackendSubEntry *) entry_ref->pending;
 	shsubent = (PgStatShared_Subscription *) entry_ref->shared_stats;
 
-- 
2.47.2