From d0fe3046eec3913829ccd5010849b25acd25fce4 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathan@postgresql.org>
Date: Wed, 30 Aug 2023 14:47:01 -0700
Subject: [PATCH] Use actual backend IDs in pg_stat_get_backend_subxact().

Unlike the other pg_stat_get_backend* functions,
pg_stat_get_backend_subxact() looks up the backend entry by using
its integer argument as a 1-based index in an internal array.  The
other functions look for the entry with the matching session
backend ID.  These numbers often match, but that isn't reliably
true.

This commit resolves this discrepancy by introducing
pgstat_get_local_beentry_by_backend_id() and using it in
pg_stat_get_backend_subxact().  We cannot use
pgstat_get_beentry_by_backend_id() because it returns a
PgBackendStatus, which lacks the locally computed additions
available in LocalPgBackendStatus that are required by
pg_stat_get_backend_subxact().

Author: Ian Barwick
Reviewed-by: Sami Imseih, Michael Paquier, Robert Haas
Discussion: https://postgr.es/m/CAB8KJ%3Dj-ACb3H4L9a_b3ZG3iCYDW5aEu3WsPAzkm2S7JzS1Few%40mail.gmail.com
Backpatch-through: 16
---
 src/backend/utils/activity/backend_status.c | 36 +++++++++++++++------
 src/backend/utils/adt/pgstatfuncs.c         |  2 +-
 src/include/utils/backend_status.h          |  1 +
 3 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/src/backend/utils/activity/backend_status.c b/src/backend/utils/activity/backend_status.c
index a4860b10fc..722c5acf38 100644
--- a/src/backend/utils/activity/backend_status.c
+++ b/src/backend/utils/activity/backend_status.c
@@ -1088,9 +1088,33 @@ cmp_lbestatus(const void *a, const void *b)
  */
 PgBackendStatus *
 pgstat_get_beentry_by_backend_id(BackendId beid)
+{
+	LocalPgBackendStatus *ret = pgstat_get_local_beentry_by_backend_id(beid);
+
+	if (ret)
+		return &ret->backendStatus;
+
+	return NULL;
+}
+
+
+/* ----------
+ * pgstat_get_local_beentry_by_backend_id() -
+ *
+ *	Like pgstat_get_beentry_by_backend_id() but with locally computed additions
+ *	(like xid and xmin values of the backend)
+ *
+ *	The beid argument is the BackendId of the desired session
+ *	(note that this is unlike pgstat_get_local_beentry_by_index()).
+ *
+ *	NB: caller is responsible for checking if the user is permitted to see this
+ *	info (especially the querystring).
+ * ----------
+ */
+LocalPgBackendStatus *
+pgstat_get_local_beentry_by_backend_id(BackendId beid)
 {
 	LocalPgBackendStatus key;
-	LocalPgBackendStatus *ret;
 
 	pgstat_read_current_status();
 
@@ -1099,14 +1123,8 @@ pgstat_get_beentry_by_backend_id(BackendId beid)
 	 * bsearch() to search it efficiently.
 	 */
 	key.backend_id = beid;
-	ret = (LocalPgBackendStatus *) bsearch(&key, localBackendStatusTable,
-										   localNumBackends,
-										   sizeof(LocalPgBackendStatus),
-										   cmp_lbestatus);
-	if (ret)
-		return &ret->backendStatus;
-
-	return NULL;
+	return bsearch(&key, localBackendStatusTable, localNumBackends,
+				   sizeof(LocalPgBackendStatus), cmp_lbestatus);
 }
 
 
diff --git a/src/backend/utils/adt/pgstatfuncs.c b/src/backend/utils/adt/pgstatfuncs.c
index 49cc887b97..dd5094a2d4 100644
--- a/src/backend/utils/adt/pgstatfuncs.c
+++ b/src/backend/utils/adt/pgstatfuncs.c
@@ -727,7 +727,7 @@ pg_stat_get_backend_subxact(PG_FUNCTION_ARGS)
 
 	BlessTupleDesc(tupdesc);
 
-	if ((local_beentry = pgstat_get_local_beentry_by_index(beid)) != NULL)
+	if ((local_beentry = pgstat_get_local_beentry_by_backend_id(beid)) != NULL)
 	{
 		/* Fill values and NULLs */
 		values[0] = Int32GetDatum(local_beentry->backend_subxact_count);
diff --git a/src/include/utils/backend_status.h b/src/include/utils/backend_status.h
index 1718ff7ce6..d51c840daf 100644
--- a/src/include/utils/backend_status.h
+++ b/src/include/utils/backend_status.h
@@ -334,6 +334,7 @@ extern uint64 pgstat_get_my_query_id(void);
  */
 extern int	pgstat_fetch_stat_numbackends(void);
 extern PgBackendStatus *pgstat_get_beentry_by_backend_id(BackendId beid);
+extern LocalPgBackendStatus *pgstat_get_local_beentry_by_backend_id(BackendId beid);
 extern LocalPgBackendStatus *pgstat_get_local_beentry_by_index(int idx);
 extern char *pgstat_clip_activity(const char *raw_activity);