Use actual backend IDs in pg_stat_get_backend_idset() and friends.

Up to now, the ID values returned by pg_stat_get_backend_idset() and
used by pg_stat_get_backend_activity() and allied functions were just
indexes into a local array of sessions seen by the last stats refresh.
This is problematic for a few reasons.  The "ID" of a session can vary
over its existence, which is surprising.  Also, while these numbers
often match the "backend ID" used for purposes like temp schema
assignment, that isn't reliably true.  We can fairly cheaply switch
things around to make these numbers actually be the sessions' backend
IDs.  The added test case illustrates that with this definition, the
temp schema used by a given session can be obtained given its PID.

While here, delete some dead code that guarded against getting
a NULL return from pgstat_fetch_stat_local_beentry().  That can't
happen as long as the caller is careful to pass an in-range array
index, as all the callers are.  (This code may not have been dead
when written, but it surely is now.)

Nathan Bossart

Discussion: https://postgr.es/m/20220815205811.GA250990@nathanxps13

Author: tglsfdc

doc/src/sgml/monitoring.sgml

@@ -5485,20 +5485,23 @@ SELECT pid, wait_event_type, wait_event FROM pg_stat_activity WHERE wait_event i
    the <structname>pg_stat_activity</structname> view, returns a set of records
    containing all the available information about each backend process.
    Sometimes it may be more convenient to obtain just a subset of this
-   information.  In such cases, an older set of per-backend statistics
+   information.  In such cases, another set of per-backend statistics
    access functions can be used; these are shown in <xref
    linkend="monitoring-stats-backend-funcs-table"/>.
-   These access functions use a backend ID number, which ranges from one
-   to the number of currently active backends.
+   These access functions use the session's backend ID number, which is a
+   small positive integer that is distinct from the backend ID of any
+   concurrent session, although a session's ID can be recycled as soon as
+   it exits.  The backend ID is used, among other things, to identify the
+   session's temporary schema if it has one.
    The function <function>pg_stat_get_backend_idset</function> provides a
-   convenient way to generate one row for each active backend for
+   convenient way to list all the active backends' ID numbers for
    invoking these functions.  For example, to show the <acronym>PID</acronym>s and
    current queries of all backends:
 
 <programlisting>
-SELECT pg_stat_get_backend_pid(s.backendid) AS pid,
-       pg_stat_get_backend_activity(s.backendid) AS query
-    FROM (SELECT pg_stat_get_backend_idset() AS backendid) AS s;
+SELECT pg_stat_get_backend_pid(backendid) AS pid,
+       pg_stat_get_backend_activity(backendid) AS query
+FROM pg_stat_get_backend_idset() AS backendid;
 </programlisting>
   </para>
 
@@ -5526,8 +5529,7 @@ SELECT pg_stat_get_backend_pid(s.backendid) AS pid,
         <returnvalue>setof integer</returnvalue>
        </para>
        <para>
-        Returns the set of currently active backend ID numbers (from 1 to the
-        number of active backends).
+        Returns the set of currently active backend ID numbers.
        </para></entry>
       </row>
 

src/backend/utils/activity/backend_status.c

@@ -846,6 +846,13 @@ pgstat_read_current_status(void)
 		/* Only valid entries get included into the local array */
 		if (localentry->backendStatus.st_procpid > 0)
 		{
+			/*
+			 * The BackendStatusArray index is exactly the BackendId of the
+			 * source backend.  Note that this means localBackendStatusTable
+			 * is in order by backend_id.  pgstat_fetch_stat_beentry() depends
+			 * on that.
+			 */
+			localentry->backend_id = i;
 			BackendIdGetTransactionIds(i,
 									   &localentry->backend_xid,
 									   &localentry->backend_xmin);
@@ -1045,26 +1052,57 @@ pgstat_get_my_query_id(void)
 	return MyBEEntry->st_query_id;
 }
 
+/* ----------
+ * cmp_lbestatus
+ *
+ *	Comparison function for bsearch() on an array of LocalPgBackendStatus.
+ *	The backend_id field is used to compare the arguments.
+ * ----------
+ */
+static int
+cmp_lbestatus(const void *a, const void *b)
+{
+	const LocalPgBackendStatus *lbestatus1 = (const LocalPgBackendStatus *) a;
+	const LocalPgBackendStatus *lbestatus2 = (const LocalPgBackendStatus *) b;
+
+	return lbestatus1->backend_id - lbestatus2->backend_id;
+}
 
 /* ----------
  * pgstat_fetch_stat_beentry() -
  *
  *	Support function for the SQL-callable pgstat* functions. Returns
- *	our local copy of the current-activity entry for one backend.
+ *	our local copy of the current-activity entry for one backend,
+ *	or NULL if the given beid doesn't identify any known session.
+ *
+ *	The beid argument is the BackendId of the desired session
+ *	(note that this is unlike pgstat_fetch_stat_local_beentry()).
  *
  *	NB: caller is responsible for a check if the user is permitted to see
  *	this info (especially the querystring).
  * ----------
  */
 PgBackendStatus *
-pgstat_fetch_stat_beentry(int beid)
+pgstat_fetch_stat_beentry(BackendId beid)
 {
+	LocalPgBackendStatus key;
+	LocalPgBackendStatus *ret;
+
 	pgstat_read_current_status();
 
-	if (beid < 1 || beid > localNumBackends)
-		return NULL;
+	/*
+	 * Since the localBackendStatusTable is in order by backend_id, we can use
+	 * 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 &localBackendStatusTable[beid - 1].backendStatus;
+	return NULL;
 }
 
 
@@ -1074,6 +1112,10 @@ pgstat_fetch_stat_beentry(int beid)
  *	Like pgstat_fetch_stat_beentry() but with locally computed additions (like
  *	xid and xmin values of the backend)
  *
+ *	The beid argument is a 1-based index in the localBackendStatusTable
+ *	(note that this is unlike pgstat_fetch_stat_beentry()).
+ *	Returns NULL if the argument is out of range (no current caller does that).
+ *
  *	NB: caller is responsible for a check if the user is permitted to see
  *	this info (especially the querystring).
  * ----------
@@ -1094,7 +1136,8 @@ pgstat_fetch_stat_local_beentry(int beid)
  * pgstat_fetch_stat_numbackends() -
  *
  *	Support function for the SQL-callable pgstat* functions. Returns
- *	the maximum current backend id.
+ *	the number of sessions known in the localBackendStatusTable, i.e.
+ *	the maximum 1-based index to pass to pgstat_fetch_stat_local_beentry().
  * ----------
  */
 int

src/backend/utils/adt/pgstatfuncs.c

@@ -415,7 +415,6 @@ pg_stat_get_backend_idset(PG_FUNCTION_ARGS)
 {
 	FuncCallContext *funcctx;
 	int		   *fctx;
-	int32		result;
 
 	/* stuff done only on the first call of the function */
 	if (SRF_IS_FIRSTCALL())
@@ -424,24 +423,33 @@ pg_stat_get_backend_idset(PG_FUNCTION_ARGS)
 		funcctx = SRF_FIRSTCALL_INIT();
 
 		fctx = MemoryContextAlloc(funcctx->multi_call_memory_ctx,
-								  2 * sizeof(int));
+								  sizeof(int));
 		funcctx->user_fctx = fctx;
 
 		fctx[0] = 0;
-		fctx[1] = pgstat_fetch_stat_numbackends();
 	}
 
 	/* stuff done on every call of the function */
 	funcctx = SRF_PERCALL_SETUP();
 	fctx = funcctx->user_fctx;
 
 	fctx[0] += 1;
-	result = fctx[0];
 
-	if (result <= fctx[1])
+	/*
+	 * We recheck pgstat_fetch_stat_numbackends() each time through, just in
+	 * case the local status data has been refreshed since we started.  It's
+	 * plenty cheap enough if not.  If a refresh does happen, we'll likely
+	 * miss or duplicate some backend IDs, but we're content not to crash.
+	 * (Refreshing midway through such a query would be problematic usage
+	 * anyway, since the backend IDs we've already returned might no longer
+	 * refer to extant sessions.)
+	 */
+	if (fctx[0] <= pgstat_fetch_stat_numbackends())
 	{
 		/* do when there is more left to send */
-		SRF_RETURN_NEXT(funcctx, Int32GetDatum(result));
+		LocalPgBackendStatus *local_beentry = pgstat_fetch_stat_local_beentry(fctx[0]);
+
+		SRF_RETURN_NEXT(funcctx, Int32GetDatum(local_beentry->backend_id));
 	}
 	else
 	{
@@ -493,17 +501,13 @@ pg_stat_get_progress_info(PG_FUNCTION_ARGS)
 		int			i;
 
 		local_beentry = pgstat_fetch_stat_local_beentry(curr_backend);
-
-		if (!local_beentry)
-			continue;
-
 		beentry = &local_beentry->backendStatus;
 
 		/*
 		 * Report values for only those backends which are running the given
 		 * command.
 		 */
-		if (!beentry || beentry->st_progress_command != cmdtype)
+		if (beentry->st_progress_command != cmdtype)
 			continue;
 
 		/* Value available to all callers */
@@ -558,24 +562,6 @@ pg_stat_get_activity(PG_FUNCTION_ARGS)
 
 		/* Get the next one in the list */
 		local_beentry = pgstat_fetch_stat_local_beentry(curr_backend);
-		if (!local_beentry)
-		{
-			int			i;
-
-			/* Ignore missing entries if looking for specific PID */
-			if (pid != -1)
-				continue;
-
-			for (i = 0; i < lengthof(nulls); i++)
-				nulls[i] = true;
-
-			nulls[5] = false;
-			values[5] = CStringGetTextDatum("<backend information not available>");
-
-			tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc, values, nulls);
-			continue;
-		}
-
 		beentry = &local_beentry->backendStatus;
 
 		/* If looking for specific PID, ignore all the others */
@@ -1180,9 +1166,9 @@ pg_stat_get_db_numbackends(PG_FUNCTION_ARGS)
 	result = 0;
 	for (beid = 1; beid <= tot_backends; beid++)
 	{
-		PgBackendStatus *beentry = pgstat_fetch_stat_beentry(beid);
+		LocalPgBackendStatus *local_beentry = pgstat_fetch_stat_local_beentry(beid);
 
-		if (beentry && beentry->st_databaseid == dbid)
+		if (local_beentry->backendStatus.st_databaseid == dbid)
 			result++;
 	}
 

src/include/utils/backend_status.h

@@ -13,6 +13,7 @@
 #include "datatype/timestamp.h"
 #include "libpq/pqcomm.h"
 #include "miscadmin.h"			/* for BackendType */
+#include "storage/backendid.h"
 #include "utils/backend_progress.h"
 
 
@@ -247,6 +248,13 @@ typedef struct LocalPgBackendStatus
 	 */
 	PgBackendStatus backendStatus;
 
+	/*
+	 * The backend ID.  For auxiliary processes, this will be set to a value
+	 * greater than MaxBackends (since auxiliary processes do not have proper
+	 * backend IDs).
+	 */
+	BackendId	backend_id;
+
 	/*
 	 * The xid of the current transaction if available, InvalidTransactionId
 	 * if not.
@@ -313,7 +321,7 @@ extern uint64 pgstat_get_my_query_id(void);
  * ----------
  */
 extern int	pgstat_fetch_stat_numbackends(void);
-extern PgBackendStatus *pgstat_fetch_stat_beentry(int beid);
+extern PgBackendStatus *pgstat_fetch_stat_beentry(BackendId beid);
 extern LocalPgBackendStatus *pgstat_fetch_stat_local_beentry(int beid);
 extern char *pgstat_clip_activity(const char *raw_activity);
 

src/test/regress/expected/stats.out

@@ -576,9 +576,9 @@ SELECT sessions > :db_stat_sessions FROM pg_stat_database WHERE datname = (SELEC
 
 -- Test pg_stat_bgwriter checkpointer-related stats, together with pg_stat_wal
 SELECT checkpoints_req AS rqst_ckpts_before FROM pg_stat_bgwriter \gset
--- Test pg_stat_wal
+-- Test pg_stat_wal (and make a temp table so our temp schema exists)
 SELECT wal_bytes AS wal_bytes_before FROM pg_stat_wal \gset
-CREATE TABLE test_stats_temp AS SELECT 17;
+CREATE TEMP TABLE test_stats_temp AS SELECT 17;
 DROP TABLE test_stats_temp;
 -- Checkpoint twice: The checkpointer reports stats after reporting completion
 -- of the checkpoint. But after a second checkpoint we'll see at least the
@@ -597,6 +597,17 @@ SELECT wal_bytes > :wal_bytes_before FROM pg_stat_wal;
  t
 (1 row)
 
+-- Test pg_stat_get_backend_idset() and some allied functions.
+-- In particular, verify that their notion of backend ID matches
+-- our temp schema index.
+SELECT (current_schemas(true))[1] = ('pg_temp_' || beid::text) AS match
+FROM pg_stat_get_backend_idset() beid
+WHERE pg_stat_get_backend_pid(beid) = pg_backend_pid();
+ match 
+-------
+ t
+(1 row)
+
 -----
 -- Test that resetting stats works for reset timestamp
 -----

src/test/regress/sql/stats.sql

@@ -303,10 +303,10 @@ SELECT sessions > :db_stat_sessions FROM pg_stat_database WHERE datname = (SELEC
 -- Test pg_stat_bgwriter checkpointer-related stats, together with pg_stat_wal
 SELECT checkpoints_req AS rqst_ckpts_before FROM pg_stat_bgwriter \gset
 
--- Test pg_stat_wal
+-- Test pg_stat_wal (and make a temp table so our temp schema exists)
 SELECT wal_bytes AS wal_bytes_before FROM pg_stat_wal \gset
 
-CREATE TABLE test_stats_temp AS SELECT 17;
+CREATE TEMP TABLE test_stats_temp AS SELECT 17;
 DROP TABLE test_stats_temp;
 
 -- Checkpoint twice: The checkpointer reports stats after reporting completion
@@ -318,6 +318,12 @@ CHECKPOINT;
 SELECT checkpoints_req > :rqst_ckpts_before FROM pg_stat_bgwriter;
 SELECT wal_bytes > :wal_bytes_before FROM pg_stat_wal;
 
+-- Test pg_stat_get_backend_idset() and some allied functions.
+-- In particular, verify that their notion of backend ID matches
+-- our temp schema index.
+SELECT (current_schemas(true))[1] = ('pg_temp_' || beid::text) AS match
+FROM pg_stat_get_backend_idset() beid
+WHERE pg_stat_get_backend_pid(beid) = pg_backend_pid();
 
 -----
 -- Test that resetting stats works for reset timestamp