From 1535930adde98162152223c1d215c1ccb0f0a9e0 Mon Sep 17 00:00:00 2001
From: Jacob Champion <jacob.champion@enterprisedb.com>
Date: Fri, 3 May 2024 15:54:58 -0700
Subject: [PATCH 1/2] pgstat: report in earlier with STATE_AUTHENTICATING

Add pgstat_bestart_pre_auth(), which reports an 'authenticating' state
while waiting for client authentication to complete. Since we hold a
transaction open across that call, and some authentication methods call
out to external systems, having a pg_stat_activity entry helps DBAs
debug when things go badly wrong.
---
 src/backend/utils/activity/backend_status.c | 37 ++++++++++++++++++---
 src/backend/utils/adt/pgstatfuncs.c         |  3 ++
 src/backend/utils/init/postinit.c           |  9 +++++
 src/include/utils/backend_status.h          |  2 ++
 4 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/src/backend/utils/activity/backend_status.c b/src/backend/utils/activity/backend_status.c
index 1ccf4c6d83..c996049bbe 100644
--- a/src/backend/utils/activity/backend_status.c
+++ b/src/backend/utils/activity/backend_status.c
@@ -71,6 +71,7 @@ static int	localNumBackends = 0;
 static MemoryContext backendStatusSnapContext;
 
 
+static void pgstat_bestart_internal(bool pre_auth);
 static void pgstat_beshutdown_hook(int code, Datum arg);
 static void pgstat_read_current_status(void);
 static void pgstat_setup_backend_status_context(void);
@@ -271,6 +272,34 @@ pgstat_beinit(void)
  */
 void
 pgstat_bestart(void)
+{
+	pgstat_bestart_internal(false);
+}
+
+
+/* ----------
+ * pgstat_bestart_pre_auth() -
+ *
+ *	Like pgstat_beinit(), above, but it's designed to be called before
+ *	authentication has been performed (so we have no user or database IDs).
+ *	Called from InitPostgres.
+ *----------
+ */
+void
+pgstat_bestart_pre_auth(void)
+{
+	pgstat_bestart_internal(true);
+}
+
+
+/* ----------
+ * pgstat_bestart_internal() -
+ *
+ *	Implementation of both flavors of pgstat_bestart().
+ *----------
+ */
+static void
+pgstat_bestart_internal(bool pre_auth)
 {
 	volatile PgBackendStatus *vbeentry = MyBEEntry;
 	PgBackendStatus lbeentry;
@@ -320,9 +349,9 @@ pgstat_bestart(void)
 	lbeentry.st_databaseid = MyDatabaseId;
 
 	/* We have userid for client-backends, wal-sender and bgworker processes */
-	if (lbeentry.st_backendType == B_BACKEND
-		|| lbeentry.st_backendType == B_WAL_SENDER
-		|| lbeentry.st_backendType == B_BG_WORKER)
+	if (!pre_auth && (lbeentry.st_backendType == B_BACKEND
+					  || lbeentry.st_backendType == B_WAL_SENDER
+					  || lbeentry.st_backendType == B_BG_WORKER))
 		lbeentry.st_userid = GetSessionUserId();
 	else
 		lbeentry.st_userid = InvalidOid;
@@ -377,7 +406,7 @@ pgstat_bestart(void)
 	lbeentry.st_gss = false;
 #endif
 
-	lbeentry.st_state = STATE_UNDEFINED;
+	lbeentry.st_state = pre_auth ? STATE_AUTHENTICATING : STATE_UNDEFINED;
 	lbeentry.st_progress_command = PROGRESS_COMMAND_INVALID;
 	lbeentry.st_progress_command_target = InvalidOid;
 	lbeentry.st_query_id = UINT64CONST(0);
diff --git a/src/backend/utils/adt/pgstatfuncs.c b/src/backend/utils/adt/pgstatfuncs.c
index 3876339ee1..f34e4a1643 100644
--- a/src/backend/utils/adt/pgstatfuncs.c
+++ b/src/backend/utils/adt/pgstatfuncs.c
@@ -366,6 +366,9 @@ pg_stat_get_activity(PG_FUNCTION_ARGS)
 
 			switch (beentry->st_state)
 			{
+				case STATE_AUTHENTICATING:
+					values[4] = CStringGetTextDatum("authenticating");
+					break;
 				case STATE_IDLE:
 					values[4] = CStringGetTextDatum("idle");
 					break;
diff --git a/src/backend/utils/init/postinit.c b/src/backend/utils/init/postinit.c
index 0805398e24..4f10e29b3d 100644
--- a/src/backend/utils/init/postinit.c
+++ b/src/backend/utils/init/postinit.c
@@ -921,6 +921,15 @@ InitPostgres(const char *in_dbname, Oid dboid,
 	{
 		/* normal multiuser case */
 		Assert(MyProcPort != NULL);
+
+		/*
+		 * Authentication can take a while, during which time we're holding a
+		 * transaction open. Fill in enough of a backend status so that DBAs can
+		 * observe what's going on. (The later call to pgstat_bestart() will
+		 * fill in the rest of the status after we've authenticated.)
+		 */
+		pgstat_bestart_pre_auth();
+
 		PerformAuthentication(MyProcPort);
 		InitializeSessionUserId(username, useroid, false);
 		/* ensure that auth_method is actually valid, aka authn_id is not NULL */
diff --git a/src/include/utils/backend_status.h b/src/include/utils/backend_status.h
index 7b7f6f59d0..f673c6a6ac 100644
--- a/src/include/utils/backend_status.h
+++ b/src/include/utils/backend_status.h
@@ -24,6 +24,7 @@
 typedef enum BackendState
 {
 	STATE_UNDEFINED,
+	STATE_AUTHENTICATING,
 	STATE_IDLE,
 	STATE_RUNNING,
 	STATE_IDLEINTRANSACTION,
@@ -309,6 +310,7 @@ extern void CreateSharedBackendStatus(void);
 
 /* Initialization functions */
 extern void pgstat_beinit(void);
+extern void pgstat_bestart_pre_auth(void);
 extern void pgstat_bestart(void);
 
 extern void pgstat_clear_backend_activity_snapshot(void);
-- 
2.34.1