From e04691d39a939ed0b8666a83dbfffce5b27c6e33 Mon Sep 17 00:00:00 2001
From: Sami Imseih <samimseih@gmail.com>
Date: Thu, 25 Jun 2026 17:26:44 -0500
Subject: [PATCH v4 3/4] pgstat: add pgstat_drop_current() for use in dshash
 scans

When iterating a stats dshash with dshash_seq_init/dshash_seq_next,
callers that want to drop entries had to use pgstat_drop_entry() which
redundantly does a dshash_find() to locate the entry that the caller
already has in hand from the scan.

Add pgstat_drop_current() which takes the PgStatShared_HashEntry and
scan status directly, performing proper pgstat bookkeeping (releasing
local entry refs, discarding pending data, and managing refcounts)
before calling pgstat_drop_entry_internal().

Also add pgstat_has_pending() to check whether an entry has unflushed
local data without doing a dshash lookup.
---
 src/backend/utils/activity/pgstat_shmem.c | 49 +++++++++++++++++++++++
 src/include/utils/pgstat_internal.h       |  3 ++
 2 files changed, 52 insertions(+)

diff --git a/src/backend/utils/activity/pgstat_shmem.c b/src/backend/utils/activity/pgstat_shmem.c
index 411d7805e80..65d53b67289 100644
--- a/src/backend/utils/activity/pgstat_shmem.c
+++ b/src/backend/utils/activity/pgstat_shmem.c
@@ -1159,6 +1159,55 @@ pgstat_drop_matching_entries(bool (*do_drop) (PgStatShared_HashEntry *, Datum),
 		pgstat_request_entry_refs_gc();
 }
 
+/*
+ * Drop a stats entry during a dshash sequential scan.
+ *
+ * The caller must have initialized the scan with exclusive locking and
+ * passes the PgStatShared_HashEntry obtained from dshash_seq_next().
+ *
+ * Returns true if the entry was freed immediately, false if other backends
+ * still hold references (in which case it is marked as dropped for later
+ * cleanup).
+ */
+bool
+pgstat_drop_current(PgStatShared_HashEntry *shent, dshash_seq_status *hstat)
+{
+	if (shent->dropped)
+		return true;
+
+	/* release local reference if we hold one */
+	if (pgStatEntryRefHash)
+	{
+		PgStat_EntryRefHashEntry *lohashent =
+			pgstat_entry_ref_hash_lookup(pgStatEntryRefHash, shent->key);
+
+		if (lohashent)
+			pgstat_release_entry_ref(lohashent->key, lohashent->entry_ref,
+									 true);
+	}
+
+	return pgstat_drop_entry_internal(shent, hstat);
+}
+
+/*
+ * Check if the given stats entry has pending local data that hasn't been
+ * flushed yet.
+ */
+bool
+pgstat_has_pending(PgStatShared_HashEntry *shent)
+{
+	if (pgStatEntryRefHash)
+	{
+		PgStat_EntryRefHashEntry *lohashent =
+			pgstat_entry_ref_hash_lookup(pgStatEntryRefHash, shent->key);
+
+		if (lohashent && lohashent->entry_ref->pending)
+			return true;
+	}
+
+	return false;
+}
+
 /*
  * Scan through all stats hash tables and drop all entries.
  */
diff --git a/src/include/utils/pgstat_internal.h b/src/include/utils/pgstat_internal.h
index d1c583a05f9..f2211eec23a 100644
--- a/src/include/utils/pgstat_internal.h
+++ b/src/include/utils/pgstat_internal.h
@@ -827,6 +827,9 @@ extern bool pgstat_drop_entry(PgStat_Kind kind, Oid dboid, uint64 objid,
 extern void pgstat_drop_all_entries(void);
 extern void pgstat_drop_matching_entries(bool (*do_drop) (PgStatShared_HashEntry *, Datum),
 										 Datum match_data);
+extern bool pgstat_drop_current(PgStatShared_HashEntry *shent,
+								dshash_seq_status *hstat);
+extern bool pgstat_has_pending(PgStatShared_HashEntry *shent);
 extern PgStat_EntryRef *pgstat_get_entry_ref_locked(PgStat_Kind kind, Oid dboid, uint64 objid,
 													bool nowait);
 extern void pgstat_reset_entry(PgStat_Kind kind, Oid dboid, uint64 objid, TimestampTz ts);
-- 
2.50.1 (Apple Git-155)