From 9f69363296058bd473efa50f8d8b995627f0dbdf Mon Sep 17 00:00:00 2001
From: Bertrand Drouvot <bertranddrouvot.pg@gmail.com>
Date: Thu, 26 Oct 2023 06:38:08 +0000
Subject: [PATCH v8] Reconcile stats in find_tabstat_entry()

---
 src/backend/utils/activity/pgstat_relation.c | 38 ++++++++++-
 src/backend/utils/adt/pgstatfuncs.c          | 66 ++------------------
 2 files changed, 41 insertions(+), 63 deletions(-)
  42.3% src/backend/utils/activity/
  57.6% src/backend/utils/adt/

diff --git a/src/backend/utils/activity/pgstat_relation.c b/src/backend/utils/activity/pgstat_relation.c
index 9876e0c1e8..4c694d925d 100644
--- a/src/backend/utils/activity/pgstat_relation.c
+++ b/src/backend/utils/activity/pgstat_relation.c
@@ -478,20 +478,52 @@ pgstat_fetch_stat_tabentry_ext(bool shared, Oid reloid)
  * Find any existing PgStat_TableStatus entry for rel_id in the current
  * database. If not found, try finding from shared tables.
  *
+ * If an entry is found, copy it and increment the copy's counters with their
+ * subtransactions counterparts. Then return the copy. The caller may need to
+ * pfree the copy (in case the MemoryContext is not reset soon after).
+ *
  * If no entry found, return NULL, don't create a new one
  */
 PgStat_TableStatus *
 find_tabstat_entry(Oid rel_id)
 {
 	PgStat_EntryRef *entry_ref;
+	PgStat_TableXactStatus *trans;
+	PgStat_TableStatus *tabentry = NULL;
+	PgStat_TableStatus *tablestatus = NULL;
 
 	entry_ref = pgstat_fetch_pending_entry(PGSTAT_KIND_RELATION, MyDatabaseId, rel_id);
 	if (!entry_ref)
+	{
 		entry_ref = pgstat_fetch_pending_entry(PGSTAT_KIND_RELATION, InvalidOid, rel_id);
+		if (!entry_ref)
+			return tablestatus;
+	}
+
+	tabentry = (PgStat_TableStatus *) entry_ref->pending;
+	tablestatus = palloc(sizeof(PgStat_TableStatus));
+	*tablestatus = *tabentry;
+
+	/*
+	 * We don't want tablestatus->trans to point to the same location as
+	 * tabentry->trans as that could lead to undesired behavior.
+	 */
+	tablestatus->trans = NULL;
+
+	/*
+	 * Live subtransactions' counts aren't in counts yet. This is not a hot
+	 * code path so it sounds ok to reconcile for tuples_inserted,
+	 * tuples_updated and tuples_deleted even if this is not what the caller
+	 * is interested in.
+	 */
+	for (trans = tabentry->trans; trans != NULL; trans = trans->upper)
+	{
+		tablestatus->counts.tuples_inserted += trans->tuples_inserted;
+		tablestatus->counts.tuples_updated += trans->tuples_updated;
+		tablestatus->counts.tuples_deleted += trans->tuples_deleted;
+	}
 
-	if (entry_ref)
-		return entry_ref->pending;
-	return NULL;
+	return tablestatus;
 }
 
 /*
diff --git a/src/backend/utils/adt/pgstatfuncs.c b/src/backend/utils/adt/pgstatfuncs.c
index 6468b6a805..998c69e241 100644
--- a/src/backend/utils/adt/pgstatfuncs.c
+++ b/src/backend/utils/adt/pgstatfuncs.c
@@ -1588,68 +1588,14 @@ PG_STAT_GET_XACT_RELENTRY_INT64(blocks_fetched)
 /* pg_stat_get_xact_blocks_hit */
 PG_STAT_GET_XACT_RELENTRY_INT64(blocks_hit)
 
-Datum
-pg_stat_get_xact_tuples_inserted(PG_FUNCTION_ARGS)
-{
-	Oid			relid = PG_GETARG_OID(0);
-	int64		result;
-	PgStat_TableStatus *tabentry;
-	PgStat_TableXactStatus *trans;
+/* pg_stat_get_xact_tuples_inserted */
+PG_STAT_GET_XACT_RELENTRY_INT64(tuples_inserted)
 
-	if ((tabentry = find_tabstat_entry(relid)) == NULL)
-		result = 0;
-	else
-	{
-		result = tabentry->counts.tuples_inserted;
-		/* live subtransactions' counts aren't in tuples_inserted yet */
-		for (trans = tabentry->trans; trans != NULL; trans = trans->upper)
-			result += trans->tuples_inserted;
-	}
+/* pg_stat_get_xact_tuples_updated */
+PG_STAT_GET_XACT_RELENTRY_INT64(tuples_updated)
 
-	PG_RETURN_INT64(result);
-}
-
-Datum
-pg_stat_get_xact_tuples_updated(PG_FUNCTION_ARGS)
-{
-	Oid			relid = PG_GETARG_OID(0);
-	int64		result;
-	PgStat_TableStatus *tabentry;
-	PgStat_TableXactStatus *trans;
-
-	if ((tabentry = find_tabstat_entry(relid)) == NULL)
-		result = 0;
-	else
-	{
-		result = tabentry->counts.tuples_updated;
-		/* live subtransactions' counts aren't in tuples_updated yet */
-		for (trans = tabentry->trans; trans != NULL; trans = trans->upper)
-			result += trans->tuples_updated;
-	}
-
-	PG_RETURN_INT64(result);
-}
-
-Datum
-pg_stat_get_xact_tuples_deleted(PG_FUNCTION_ARGS)
-{
-	Oid			relid = PG_GETARG_OID(0);
-	int64		result;
-	PgStat_TableStatus *tabentry;
-	PgStat_TableXactStatus *trans;
-
-	if ((tabentry = find_tabstat_entry(relid)) == NULL)
-		result = 0;
-	else
-	{
-		result = tabentry->counts.tuples_deleted;
-		/* live subtransactions' counts aren't in tuples_deleted yet */
-		for (trans = tabentry->trans; trans != NULL; trans = trans->upper)
-			result += trans->tuples_deleted;
-	}
-
-	PG_RETURN_INT64(result);
-}
+/* pg_stat_get_xact_tuples_deleted */
+PG_STAT_GET_XACT_RELENTRY_INT64(tuples_deleted)
 
 Datum
 pg_stat_get_xact_function_calls(PG_FUNCTION_ARGS)
-- 
2.34.1